SpringBoot2.x整合Mongo实现多数据库切换

前言

本项目Mongo的多数据库切换的核心其实还是利用AOP技术以及自定义的MongoTemplate。切点就是每一个dao操作类，大致思路就是在进行操作之前构造对应数据库的MongoTemplate，这样就可以利用构造的MongoaTemplate进行操作了。

MultiMongoTemplate

首先是自定义的MongoTemplate，这个不用多讲，就是对mongo执行各种操作的关键所在。这里多数据库的切换，说白了就是在需要切换的时候，利用MongoDbFactory构造对应的MongoTemplate即可。然后保存到ThreadLocal中，每次获取MongoDatabase的时候从里面取即可

Component(value = "mongoTemplate")
@Slf4j
public class MultiMongoTemplate extends MongoTemplate {

    private static volatile ThreadLocal<MongoDbFactory> mongoDbFactoryThreadLocal;

    public MultiMongoTemplate(MongoDbFactory mongoDbFactory) {
        super(mongoDbFactory);
        if (mongoDbFactoryThreadLocal == null) {
            synchronized (MultiMongoTemplate.class) {
                if (mongoDbFactoryThreadLocal == null) {
                    mongoDbFactoryThreadLocal = new ThreadLocal<>();
                    mongoDbFactoryThreadLocal.set(mongoDbFactory);
                }
            }
        }
    }
    public void setMongoDbFactory(MongoDbFactory factory){
        mongoDbFactoryThreadLocal.set(factory);
    }

    public void removeMongoDbFactory(){
        mongoDbFactoryThreadLocal.remove();
    }

    @Override
    public MongoDatabase getDb() {
        return mongoDbFactoryThreadLocal.get().getDb();
    }

    @Override
    protected MongoDatabase doGetDatabase(){
        return mongoDbFactoryThreadLocal.get().getDb();
    }
}

数据库动态切换

@Component
@Aspect
@Slf4j
public class MongoSwitchAspect {

    @Value("${spring.data.mongodb.uri}")
    private String mongoUri;

    private Map<String, MongoDbFactory> mongoDbFactoryMap = new ConcurrentHashMap<>();

    // 切点除了MongoTemplate的set方法
    @Pointcut("execution(* com.itoyoung.dao.mongo.*.*(..)) && !execution(* com.itoyoung.dao.mongo.*.set*(..))")
    public void routeMongoDB() {
    }

    @Around("routeMongoDB()")
    public Object routeMongoDB(ProceedingJoinPoint joinPoint) {

        // 判断数据库是否和之前的一致，一致就不切换
        if (DynamicMongoDBContextHolder.switchDB()) {

            //获取需要访问的项目数据库
            String dbName = DynamicMongoDBContextHolder.getDBRouterKey();
            try {

                Object o = joinPoint.getTarget();
                Class<?> clazz = o.getClass();

                Class<?> superclass = clazz.getSuperclass();
                if (superclass != BaseMongoDao.class) {
                    log.error("MongoSwitch: class: {} 未继承BaseMongoDao或非Mongo操作类，请勿放入dao.mongo路径下",
                            clazz.getName());
                    throw new BaseRuntimeException("数据库切换异常");
                }

                Field mongoTemplateField = superclass.getDeclaredField("mongoTemplate");
                if (null == mongoTemplateField) {
                    log.error("MongoSwitch: class: {} 不存在mongoTemplate", clazz.getName());
                    throw new BaseRuntimeException("数据库切换异常");
                }
                mongoTemplateField.setAccessible(true);

                MultiMongoTemplate mongoTemplate = (MultiMongoTemplate) mongoTemplateField.get(o);

                //设置MongoFactory
                if (null == mongoTemplate) {
                    mongoTemplate = new MultiMongoTemplate(getDbFactory(dbName));
                } else {
                    mongoTemplate.setMongoDbFactory(getDbFactory(dbName));
                }
                //重新赋值
                mongoTemplateField.set(o, mongoTemplate);
            } catch (Exception e) {
                log.error("MongoSwitch: DB: {} 切换数据库失败: {}", dbName, e);
            }
        }

        Object result = null;
        try {
            result = joinPoint.proceed();
//            if (StringUtils.isBlank(switchOrg)) {
//                //清理ThreadLocal的变量
//                mongoTemplate.removeMongoDbFactory();
//            }

        } catch (Throwable t) {
            log.error("MongoSwitch: DB: {} 方法执行失败: {}", DynamicMongoDBContextHolder.getDBRouterKey(), t);
        }
        return result;
    }

    /**
     * 根据需要切换的数据库获取MongoDbFactory
     *
     * @param dbName
     * @return
     */
    private SimpleMongoDbFactory getDbFactory(String dbName) {
        //查找项目对应的MongFactory
        SimpleMongoDbFactory simpleMongoDbFactory = (SimpleMongoDbFactory) mongoDbFactoryMap.get(dbName);

        //实例化
        if (simpleMongoDbFactory == null) {
            MongoClient mongoClient = new MongoClient(new MongoClientURI(mongoUri));

            simpleMongoDbFactory = new SimpleMongoDbFactory(mongoClient, dbName);

            mongoDbFactoryMap.put(dbName, simpleMongoDbFactory);
        }
        return simpleMongoDbFactory;
    }

}

动态切换操作类

协助在业务代码中切换mongo数据库

@Slf4j
public class DynamicMongoDBContextHolder {

    /**
     * 线程级别的私有变量(存储需切换的组织)
     */
    private static final ThreadLocal<String> MONGO_DB_HOLDER = new ThreadLocal<>();

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

    /**
     * 获取当前操作的数据库
     *
     * @return
     */
    public static String getDBRouterKey() {
        return MONGO_DB_HOLDER.get();
    }

    /**
     * 切换mongo库
     *
     * @param dbName
     */
    public static void setDBRouterKey(String dbName) {
        if (StringUtils.isEmpty(dbName)) {
            log.info("要切换的mongo库为空");
            throw new BaseRuntimeException("当前切换的数据源不存在");
        }
        // 将切换之前的数据库保存
        String beforeDB = MONGO_DB_HOLDER.get();
        if (StringUtils.isNotEmpty(beforeDB) && !beforeDB.equals(dbName)) {
            BEFORE_DB__HOLDER.set(beforeDB);
        }
        //将需切的mongo库存到当前线程的threadLocal内
        MONGO_DB_HOLDER.set(dbName);
    }

    /**
     * 切换请求最开始的数据库
     */
    public static void setRootDBRouterKey() {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        if (requestAttributes == null) {
            throw new BaseRuntimeException("非Request请求，切换Root数据源失败");
        }

        String dbName = ((ServletRequestAttributes) requestAttributes).getRequest().getParameter("dbName");
        setDBRouterKey(dbName);
    }

    /**
     * 判断是否需要切换数据库
     *
     * @return
     */
    public static boolean switchDB() {
        String nowDB = MONGO_DB_HOLDER.get();
        if (StringUtils.isEmpty(nowDB)) {
            throw new BaseRuntimeException("当前数据源不存在");
        }
        // 之前的集合和当前的集合不一致 就需要切换
        return !nowDB.equals(BEFORE_DB__HOLDER.get());
    }

    /**
     * 移除
     */
    public static void remove() {
        MONGO_DB_HOLDER.remove();
        BEFORE_DB__HOLDER.remove();
    }
}

Mongo操作基础类

之前整合Mongo的那篇有写，但是需要做出一点修改

@Slf4j
public abstract class BaseMongoDao<T> {

    /**
     * 反射获取泛型类型
     *
     * @return
     */
    protected abstract Class<T> getEntityClass();

    @Resource
    public MultiMongoTemplate mongoTemplate;

    public MultiMongoTemplate getMongoTemplate() {
        return mongoTemplate;
    }

    public void setMongoTemplate(MultiMongoTemplate mongoTemplate) {
        this.mongoTemplate = mongoTemplate;
    }

    /***
     * 保存一个对象
     * @param t
     */
    public void saveWithColletion(T t, String collection) {
        log.info("-------------->MongoDB save start");
        this.mongoTemplate.save(t, collection);
    }

    /**
    * 根据条件查询集合
    *
    * @param object
    * @return
    */
    public List<T> queryListWithCollection(T object, String collection) {
        Query query = getQueryByObject(object);
        log.info("-------------->MongoDB find start");
        return mongoTemplate.find(query, this.getEntityClass(), collection);
    }
    // 省略...
}

拦截器

在请求进来时进行拦截，切换到对应的数据库

/**
 * 拦截器配置类
 * 
 * @author ：itoyoung
 * @date ：2019-06-02 17:23
 * 
 */
@Configuration
public class WebConfigurer implements WebMvcConfigurer {

    @Resource
    private HttpRequestInterceptor httpRequestInterceptor;

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        //将自定义的拦截器加入配置中。
        registry.addInterceptor(httpRequestInterceptor).addPathPatterns("/**");
    }
}

@Component
public class HttpRequestInterceptor extends HandlerInterceptorAdapter {

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String dbName = request.getParameter("dbName");
        if (StringUtil.isEmpty(dbName)) {
            // 当没有传mongo数据库时，给一个默认的数据库
            dbName = "xzjj";
        } 
        DynamicMongoDBContextHolder.setDBRouterKey(dbName);
        return super.preHandle(request, response, handler);
    }
}

简单应用

/**
 * 用户日志mongo服务
 *
 * @author : itoyoung
 * @date :   2019-06-02 18:06
 */
@Service
public class UserLogMongoServiceImpl implements IUserLogMongoService {

    private final static String USER_LOG = "user_log";

    @Resource
    private UserLogMongoDao userLogMongoDao;

    @Override
    public void insert(MongoUserLog userLog) {
        // 请求的数据库
        userLogMongoDao.saveWithColletion(userLog, USER_LOG);
        // 手动切换数据库
        DynamicMongoDBContextHolder.setDBRouterKey("yjdf");
        userLogMongoDao.saveWithColletion(userLog, USER_LOG);
    }


    @Override
    public List<MongoUserLog> listByQuery(MongoUserLog query) {
        return userLogMongoDao.queryListWithCollection(query, USER_LOG);
    }
}