MongoDB数据源动态切换
项目需要根据业务来分库，比如任务A相关数据入库到MongoDB-A，任务B相关数据入库到MongoDB-B；
网上搜索了下大概方案都是依赖AOP来实现，代码也都千篇一律，BUG百出；于是修改了下，大概的思路如下：
- 切面放在了MongoTemplate中
@Aspect
@Component
@Slf4j
public class MongodbSpecialSourceAspect {
private final String POINT_CUT = "execution (* org.springframework.data.mongodb.core.MongoTemplate.*(..))";
@Around(value = POINT_CUT)
public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
Object o = joinPoint.getTarget();
CustomMongoTemplate customMongoTemplate = null;
// 判断是否是自定义的MongoTemplate
if (o.getClass() == CustomMongoTemplate.class) {
customMongoTemplate = (CustomMongoTemplate) o;
// TODO 判断是否切换逻辑
if(xxxx) {
customMongoTemplate.switchMongoDbFactory(CustomMongoTemplate.EDU_DB_KEY);
}
}
Object result = joinPoint.proceed();
if (customMongoTemplate != null) {
customMongoTemplate.restoreMaster();
}
return result;
}
}
- 声明MongoTemplate的子类，重写getMongoDbFactory方法
@Slf4j
public class CustomMongoTemplate extends MongoTemplate {
// 存放当前的数据源dbFactory
private ThreadLocal<MongoDbFactory> mongoDbFactoryThreadLocal;
// 存放初始化多数据源的dbFactory
private final Map<Integer, MongoDbFactory> mongoDbFactoriesMap;
public final static int MASTER_DB_KEY = 0;
public final static int EDU_DB_KEY = 1;
public CustomMongoTemplate(Map<Integer, MongoDbFactory> mongoDbFactoriesMap, MongoDbFactory mongoDbFactory){
super(mongoDbFactory);
if (!mongoDbFactoriesMap.containsKey(0)) {
throw new BeanInitializationException("mongoDbFactoriesMap must contain an entry with key=0.");
}
this.mongoDbFactoriesMap = mongoDbFactoriesMap;
if(mongoDbFactoryThreadLocal==null) {
mongoDbFactoryThreadLocal = new ThreadLocal<>();
}
}
public void switchMongoDbFactory(Integer mongoDbFactoryKey){
if (!mongoDbFactoriesMap.containsKey(mongoDbFactoryKey)) {
throw new IllegalArgumentException("The entry of the current key cannot be found in mongoDbFactoriesMap.");
}
mongoDbFactoryThreadLocal.set(mongoDbFactoriesMap.get(mongoDbFactoryKey));
}
public void restoreMaster() {
mongoDbFactoryThreadLocal.set(mongoDbFactoriesMap.get(MASTER_DB_KEY));
}
///////////////////////
@Override
public MongoDbFactory getMongoDbFactory() {
if (mongoDbFactoryThreadLocal.get() == null) {
restoreMaster();
}
return mongoDbFactoryThreadLocal.get();
}
@Override
protected MongoDatabase doGetDatabase() {
return MongoDatabaseUtils.getDatabase(getMongoDbFactory(), SessionSynchronization.ON_ACTUAL_TRANSACTION);
}
@Override
public SessionScoped withSession(ClientSessionOptions options) {
Assert.notNull(options, "ClientSessionOptions must not be null!");
return withSession(() -> getMongoDbFactory().getSession(options));
}
////////////////////
主要是重写getMongoDbFactory方法来实时改变dbFactory，因为可以查阅源码看到所有的dbFactory都是来源于此方法。当然也有例外，比如
doGetDatabase和withSession还有findDistinct方法，所以我们重写了doGetDatabase和withSession方法，由直接使用私有变量mongoDbFactory改为调用getMongoDbFactory方法来获取mongoDbFactory。
而findDistinct方法中使用了大量的私有变量，导致无法重写此方法；如果有用到，可以重新实现下此方法；
- 声明MongoTemplate的Bean实例
@Configuration
@ConfigurationProperties(prefix = "spring.data.mongodb.file")
public class FileMongoConfig {
private String url;
private String eduUrl;
@Bean(name = "fileMongoTemplate")
public MongoTemplate getMongoTemplate() {
// 将需要切换的数据源都放到map中
Map<Integer, MongoDbFactory> mongoDbFactoriesMap = new HashMap<>(1<<3);
MongoDbFactory mongoDbFactory = this.fileMongoDbFactory();
mongoDbFactoriesMap.put(CustomMongoTemplate.MASTER_DB_KEY, mongoDbFactory);
mongoDbFactoriesMap.put(CustomMongoTemplate.EDU_DB_KEY, this.fileEduMongoDbFactory());
// 这里返回我们自定义的MongoTemplate
return new CustomMongoTemplate(mongoDbFactoriesMap, mongoDbFactory);
}
public MongoDbFactory fileMongoDbFactory() {
return new SimpleMongoDbFactory(new MongoClientURI(getUrl()));
}
public MongoDbFactory fileEduMongoDbFactory() {
return new SimpleMongoDbFactory(new MongoClientURI(getEduUrl()));
}
public String getUrl() {
return url;
}
public void setUrl(String url) {
this.url = url;
}
public String getEduUrl() {
return eduUrl;
}
public void setEduUrl(String eduUrl) {
this.eduUrl = eduUrl;
}
}
-
以上三步后，我们准备工作就完成了；此时进行一个mongo查询时，切面拦截，判断是否是自定义的MongoTemplate，如果是则进行逻辑判断是否切换数据源；
在执行完mongo语句后，切换回主数据源；完成数据源的切换；
缺陷
- 此方案缺陷很明显，aop拿不到方法参数，如果场景依赖的是方法的传参来判断是否切换数据源，那么用aop就无法完成；此方案只适合整体切换的场景；
转换思路，如果我需要依赖collectionName来切换数据源或者部分入参来切换数据源，该怎么实现呢？
- 首先aop肯定是拿不到参数的，就算使用threadlocal传输也很难实现
- 那么我们直接重写MongoTemplate的方法，在所有方法前都进行逻辑判断是否切换，在方法执行后还原
- 删除aspect类
- 修改下CustomMongoTemplate，我们仔细筛选MongoTemplate的所有参数有collectionName的方法，大部分底层调用都是一样的。那么我们只需要重写这些底层的方法就可以；最终如下：
@Slf4j
public class CustomMongoTemplate extends MongoTemplate {
private ThreadLocal<MongoDbFactory> mongoDbFactoryThreadLocal;
// 用于分区map, 0: master-db 1: edu-db
private final Map<Integer, MongoDbFactory> mongoDbFactoriesMap;
public final static int MASTER_DB_KEY = 0;
public final static int EDU_DB_KEY = 1;
public CustomMongoTemplate(Map<Integer, MongoDbFactory> mongoDbFactoriesMap, MongoDbFactory mongoDbFactory){
super(mongoDbFactory);
if (!mongoDbFactoriesMap.containsKey(0)) {
throw new BeanInitializationException("mongoDbFactoriesMap must contain an entry with key=0.");
}
this.mongoDbFactoriesMap = mongoDbFactoriesMap;
if(mongoDbFactoryThreadLocal==null) {
mongoDbFactoryThreadLocal = new ThreadLocal<>();
}
}
public CustomMongoTemplate(Map<Integer, MongoDbFactory> mongoDbFactoriesMap, MongoDbFactory mongoDbFactory, @Nullable MongoConverter mongoConverter){
super(mongoDbFactory, mongoConverter);
if (!mongoDbFactoriesMap.containsKey(0)) {
throw new BeanInitializationException("mongoDbFactoriesMap must contain an entry with key=0.");
}
this.mongoDbFactoriesMap = mongoDbFactoriesMap;
if(mongoDbFactoryThreadLocal==null) {
mongoDbFactoryThreadLocal = new ThreadLocal<>();
}
}
public void switchMongoDbFactory(Integer mongoDbFactoryKey){
if (!mongoDbFactoriesMap.containsKey(mongoDbFactoryKey)) {
throw new IllegalArgumentException("The entry of the current key cannot be found in mongoDbFactoriesMap.");
}
mongoDbFactoryThreadLocal.set(mongoDbFactoriesMap.get(mongoDbFactoryKey));
}
public void restoreMaster() {
mongoDbFactoryThreadLocal.set(mongoDbFactoriesMap.get(MASTER_DB_KEY));
}
public void checkAndSwitchDbFactory(String collectionName) {
// TODO 判断业务相关逻辑，切换到相关的DbFactory
this.switchMongoDbFactory(EDU_DB_KEY);
}
///////////////////////
@Override
public MongoDbFactory getMongoDbFactory() {
if (mongoDbFactoryThreadLocal.get() == null) {
restoreMaster();
}
return mongoDbFactoryThreadLocal.get();
}
@Override
protected MongoDatabase doGetDatabase() {
return MongoDatabaseUtils.getDatabase(getMongoDbFactory(), SessionSynchronization.ON_ACTUAL_TRANSACTION);
}
@Override
public SessionScoped withSession(ClientSessionOptions options) {
Assert.notNull(options, "ClientSessionOptions must not be null!");
return withSession(() -> getMongoDbFactory().getSession(options));
}
////////////////////
@Override
public <S, T> List<T> doFind(String collectionName, Document query, Document fields, Class<S> sourceClass, Class<T> targetClass, CursorPreparer preparer) {
checkAndSwitchDbFactory(collectionName);
List<T> ts = super.doFind(collectionName, query, fields, sourceClass, targetClass, preparer);
restoreMaster();
return ts;
}
@Override
protected void executeQuery(Query query, String collectionName, DocumentCallbackHandler documentCallbackHandler, CursorPreparer preparer) {
checkAndSwitchDbFactory(collectionName);
super.executeQuery(query, collectionName, documentCallbackHandler, preparer);
restoreMaster();
}
@Override
public <T> T execute(String collectionName, CollectionCallback<T> callback) {
checkAndSwitchDbFactory(collectionName);
T x = super.execute(collectionName, callback);
restoreMaster();
return x;
}
@Override
public MongoCollection<Document> getCollection(String collectionName) {
checkAndSwitchDbFactory(collectionName);
MongoCollection<Document> x = super.getCollection(collectionName);
restoreMaster();
return x;
}
@Override
public boolean collectionExists(String collectionName) {
checkAndSwitchDbFactory(collectionName);
boolean x = super.collectionExists(collectionName);
restoreMaster();
return x;
}
@Override
public <T> List<T> findDistinct(Query query, String field, String collectionName, Class<?> entityClass, Class<T> resultClass) {
checkAndSwitchDbFactory(collectionName);
List<T> x = super.findDistinct(query, field, collectionName, entityClass, resultClass);
restoreMaster();
return x;
}
@Override
public <T> GeoResults<T> geoNear(NearQuery near, Class<?> domainType, String collectionName, Class<T> returnType) {
checkAndSwitchDbFactory(collectionName);
GeoResults<T> x = super.geoNear(near, domainType, collectionName, returnType);
restoreMaster();
return x;
}
@Override
protected <T> T doInsert(String collectionName, T objectToSave, MongoWriter<T> writer) {
checkAndSwitchDbFactory(collectionName);
T x = super.doInsert(collectionName, objectToSave, writer);
restoreMaster();
return x;
}
@Override
protected <T> Collection<T> doInsertBatch(String collectionName, Collection<? extends T> batchToSave, MongoWriter<T> writer) {
checkAndSwitchDbFactory(collectionName);
Collection<T> x = super.doInsertBatch(collectionName, batchToSave, writer);
restoreMaster();
return x;
}
@Override
protected <T> T doSave(String collectionName, T objectToSave, MongoWriter<T> writer) {
checkAndSwitchDbFactory(collectionName);
T x = super.doSave(collectionName, objectToSave, writer);
restoreMaster();
return x;
}
@Override
public <T> List<T> findAll(Class<T> entityClass, String collectionName) {
checkAndSwitchDbFactory(collectionName);
List<T> x = super.findAll(entityClass, collectionName);
restoreMaster();
return x;
}
@Override
public <T> MapReduceResults<T> mapReduce(Query query, String inputCollectionName, String mapFunction, String reduceFunction, MapReduceOptions mapReduceOptions, Class<T> entityClass) {
checkAndSwitchDbFactory(inputCollectionName);
MapReduceResults<T> x = super.mapReduce(query, inputCollectionName, mapFunction, reduceFunction, mapReduceOptions, entityClass);
restoreMaster();
return x;
}
@Override
public <T> List<T> mapReduce(Query query, Class<?> domainType, String inputCollectionName, String mapFunction, String reduceFunction, MapReduceOptions mapReduceOptions, Class<T> resultType) {
checkAndSwitchDbFactory(inputCollectionName);
List<T> x = super.mapReduce(query, domainType, inputCollectionName, mapFunction, reduceFunction, mapReduceOptions, resultType);
restoreMaster();
return x;
}
@Override
public <T> GroupByResults<T> group(Criteria criteria, String inputCollectionName, GroupBy groupBy, Class<T> entityClass) {
checkAndSwitchDbFactory(inputCollectionName);
GroupByResults<T> x = super.group(criteria, inputCollectionName, groupBy, entityClass);
restoreMaster();
return x;
}
@Override
protected <O> AggregationResults<O> doAggregate(Aggregation aggregation, String collectionName, Class<O> outputType, AggregationOperationContext context) {
checkAndSwitchDbFactory(collectionName);
AggregationResults<O> x = super.doAggregate(aggregation, collectionName, outputType, context);
restoreMaster();
return x;
}
@Override
protected MongoCollection<Document> doCreateCollection(String collectionName, Document collectionOptions) {
checkAndSwitchDbFactory(collectionName);
MongoCollection<Document> x = super.doCreateCollection(collectionName, collectionOptions);
restoreMaster();
return x;
}
@Override
protected <T> T doFindOne(String collectionName, Document query, Document fields, Class<T> entityClass) {
checkAndSwitchDbFactory(collectionName);
T x = super.doFindOne(collectionName, query, fields, entityClass);
restoreMaster();
return x;
}
@Override
protected <S, T> List<T> doFind(String collectionName, Document query, Document fields, Class<S> entityClass, CursorPreparer preparer, DocumentCallback<T> objectCallback) {
checkAndSwitchDbFactory(collectionName);
List<T> x = super.doFind(collectionName, query, fields, entityClass, preparer, objectCallback);
restoreMaster();
return x;
}
@Override
protected <T> T doFindAndRemove(String collectionName, Document query, Document fields, Document sort, Collation collation, Class<T> entityClass) {
checkAndSwitchDbFactory(collectionName);
T x = super.doFindAndRemove(collectionName, query, fields, sort, collation, entityClass);
restoreMaster();
return x;
}
@Override
protected <T> T doFindAndModify(String collectionName, Document query, Document fields, Document sort, Class<T> entityClass, Update update, FindAndModifyOptions options) {
checkAndSwitchDbFactory(collectionName);
T x = super.doFindAndModify(collectionName, query, fields, sort, entityClass, update, options);
restoreMaster();
return x;
}
@Override
protected <T> T doFindAndReplace(String collectionName, Document mappedQuery, Document mappedFields, Document mappedSort, com.mongodb.client.model.Collation collation, Class<?> entityType, Document replacement, FindAndReplaceOptions options, Class<T> resultType) {
checkAndSwitchDbFactory(collectionName);
T x = super.doFindAndReplace(collectionName, mappedQuery, mappedFields, mappedSort, collation, entityType, replacement, options, resultType);
restoreMaster();
return x;
}
}
收工；
时在中春,阳和方起
