基于注解动态切换Mongo数据源

定义MongoContext

pom依赖

<dependency>  
    <groupId>org.springframework.boot</groupId>  
    <artifactId>spring-boot-starter-data-mongodb</artifactId>  
</dependency>

import com.mongodb.client.MongoClient;
​
import java.util.HashMap;
import java.util.Map;
​
/**
 * @author Zpain
 * @date 2023/4/21 15:12
 */
public class MongoContext {
    //用来存放所有mongoClient
    private static final Map<String, MongoClient> MONGO_CLIENT_DB_MAP = new HashMap<>();
    //当前线程使用的mongoClient
    private static final ThreadLocal<MongoClient> MONGO_DB_THREAD_LOCAL = new ThreadLocal<>();
​
    public static void putMongoDb(String name, MongoClient factory) {
        MONGO_CLIENT_DB_MAP.put(name, factory);
    }
​
    public static Map<String, MongoClient> getMongoClient() {
        return MONGO_CLIENT_DB_MAP;
    }
​
    public static void setMongoDbClient(String name) {
        MONGO_DB_THREAD_LOCAL.set(MONGO_CLIENT_DB_MAP.get(name));
    }
​
    public static MongoClient getMongoDbClient() {
        return MONGO_DB_THREAD_LOCAL.get();
    }
​
    
    public static void removeMongoDbClient() {
        MONGO_DB_THREAD_LOCAL.remove();
    }
​
}

注入mongo到MongoContext中

public static void main(String[] args) {
        try {
            String id = "";
            //校验当前id是否已经注入
            Set<String> strings = MongoContext.getMongoClient().keySet();
            boolean b = false;
            if (!CollectionUtils.isEmpty(strings)) {
                b = strings.contains(id);
            }
            if (!b) {
                String dbPassword ="";
                String user = "";
                String url = "";
                String port = "";
                String db = "";
                String uri = String.format("mongodb://%s:%s@%s:%s/%s", user, dbPassword
                        , url, port, db);
​
                MongoCredential credential = MongoCredential.createCredential(user, db, dbPassword.toCharArray());
                MongoClientSettings settings = MongoClientSettings.builder()
                        .applyConnectionString(new ConnectionString(uri))
                        .applyToClusterSettings(builder -> {
                            builder.serverSelectionTimeout(2, TimeUnit.SECONDS);
                        })
                        .credential(credential)
                        .build();
                MongoClient client = MongoClients.create(settings);
                //检测mongo链接是否正确
                ListDatabasesIterable<Document> documents = client.listDatabases();
                for (Document d : documents) {
                    d.get("test");
                }
                //存放进mongoContext中
                MongoContext.putMongoDb(id, client);
            }
        } catch (Exception e) {
            log.error("error:{}", e.getMessage());
        }
    }

定义自定义切换注解

import java.lang.annotation.*;
​
/**
 * @author Zpain
 * @date 2023/4/21 15:20
 */
@Target({ElementType.METHOD})
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface MongoSwitch {
}

通过aop去切换数据源

​
import com.mongodb.client.MongoClient;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.After;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.stereotype.Component;
​
import java.lang.reflect.Field;
​
/**
 * @author Zpain
 * @date 2023/4/21 15:21
 */
@Component
@Slf4j
@Aspect
public class MongoAspect {
​
    @Pointcut("@annotation(com.xxx.MongoSwitch)")
    public void mongoSwitch() {
​
    }
​
    @Before("mongoSwitch()")
    public void before(JoinPoint point) {
        try {
            Object[] args = point.getArgs();
            for (Object obj : args) {
                //mongo在数据库中配置的id
                Long id = (Long) obj;
                //1.自行获取mongo配置数据
                Database database = xxx.getDatabase(id);
                //2.注入到mongoContext中
                dynamicDataSourceService.setMongo(database);
                //切换到准备使用的mongo
                MongoContext.setMongoDbClient(id.toString());
                MongoClient client = MongoContext.getMongoDbClient();
                //dbname mongo库名称
                MongoTemplate mongoTemplate = new MongoTemplate(client, dbName);
                //将mongoClient注入到Spring的mongoTemplate中 
                Object targetObject = point.getTarget();
                //给mongoTemplate赋值
                Field template = targetObject.getClass().getDeclaredField("mongoTemplate");
                if (template.getType().equals(MongoTemplate.class)) {
                    try {
                        template.setAccessible(true);
                        template.set(targetObject, mongoTemplate);
                    } catch (IllegalAccessException e) {
                        throw new RuntimeException(String.format("Error injecting MongoTemplate instance into %s", targetObject.getClass().getName()), e);
                    }
                    return;
                }
            }
        } catch (Exception e) {
            log.error("MongoAspect[before] error:{}", e.getMessage());
            throw new BusinessException(ExceptionConstants.MONGO_ERROR.getMessage());
        }
    }
​
    /*
    * 结束之后释放 否则相同线程进来可能无法切换到正确的数据源
    */
    @After("mongoSwitch()")
    public void after(JoinPoint point) {
        try {
            MongoContext.removeMongoDbClient();
        } catch (Exception e) {
            log.error("MongoAspect[after] error:{}", e.getMessage());
        }
    }