springboot使用cglib代理实现多源数据库动态切换

需求场景

想要根据业务的一些参数和场景动态切换不同的mongo数据库,比如根据用户请求http header里面的参数来使用不同的数据库。其他数据库可以参照这个写的去动态实现。

实现思路

所有的数据库操作都用同一个MongoTemplate，使用cglib代理这个客户端，根据不同的参数，切换不通的客户端来进行最终调用。

代码实现

这里直接上代码，简单粗暴

简单对MongoTemplate进行一层包装

创建一个MyMongoTemplate类来做一个壳

public class MyMongoTemplate extends MongoTemplate {
    public MyMongoTemplate() {
        super(
                new SimpleMongoClientDatabaseFactory(
                        new ConnectionString("mongodb://name:pass@localhost:27017/test")));
    }

    public MyMongoTemplate(MongoDatabaseFactory mongoDbFactory, MongoConverter mongoConverter) {
        super(mongoDbFactory, mongoConverter);
    }
}

其中，数据库连接可以随便写一个，或者使用一个默认的连接，这个连接实际上是用不到的。因为MongoTemplate不是无参构造，所以这里还需要实现一个带参的构造函数，这里直接super就可以了。

编写代理逻辑

 public static MyMongoTemplate getTmp(MongoMappingContext context, String uri) {
        MappingMongoConverter converter =
                new MappingMongoConverter(new DefaultDbRefResolver(mongoDbFactory(uri)), context);
        converter.setTypeMapper(new DefaultMongoTypeMapper(null));

        return new MyMongoTemplate(mongoDbFactory(uri), converter);
    }

    public static Object createProxy(String usUri, String idUri, MongoMappingContext context) {
        MyMongoTemplate mongoUs = getTmp(context, usUri);
        MyMongoTemplate mongoId = getTmp(context, idUri);
        Enhancer enhancer = new Enhancer();
        enhancer.setSuperclass(MyMongoTemplate.class);
        enhancer.setCallback(
                new MethodInterceptor() {
                    @Override
                    public Object intercept(
                            Object obj, Method method, Object[] args, MethodProxy proxy)
                            throws Exception {
                        long mongoStart = System.currentTimeMillis();
                        try {
                            if (method.getName().equals("toString")) {
                                return proxy.getClass().getName()
                                        + "@"
                                        + Integer.toHexString(System.identityHashCode(proxy));
                            }
                            HttpServletRequest request = getCurrentHttpRequest();
                            if (null == request) {
                                return method.invoke(mongoUs, args);
                            }
                            String region = request.getHeader("site-region");
                            if (!StringUtils.hasLength(region) || "US".equals(region)) {
                                return method.invoke(mongoUs, args);
                            }
                            if ("ID".equals(region)) {
                                return method.invoke(mongoId, args);
                            } else {
                                return method.invoke(mongoUs, args);
                            }
                        } catch (Exception e) {
                            log.error("mongo proxy error,e", e);
                            throw e;
                        } finally {
                            log.info("mongodb cost:{}", System.currentTimeMillis() - mongoStart);
                        }
                    }
                });
        return enhancer.create();
    }

这里createProxy方法接收两个mongodb uri,然后会返回一个代理的对象。

配置参数与创建代理

这里直接创建一个Bean,返回一个MongoTemplate对象，注入到spring里面。后续只需要使用这个MongoTemplate对象，就可以根据http header的不同动态切换不同的数据源来处理了

 @Bean(name = "baseTemplate")
    @Primary
    public MongoTemplate initBase(
            MongoMappingContext context,
            @Value("${spring.data.mongodb.base.uri.us}") String uriUs,
            @Value("${spring.data.mongodb.base.uri.id}") String uriId) {
        return (MyMongoTemplate) MongoProxy.createProxy(uriUs, uriId, context);

代码中会用到的import

import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Primary;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
import com.mongodb.ConnectionString;
import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cglib.proxy.Enhancer;
import org.springframework.cglib.proxy.MethodInterceptor;
import org.springframework.cglib.proxy.MethodProxy;
import org.springframework.data.mongodb.MongoDatabaseFactory;
import org.springframework.data.mongodb.core.SimpleMongoClientDatabaseFactory;
import org.springframework.data.mongodb.core.convert.DefaultDbRefResolver;
import org.springframework.data.mongodb.core.convert.DefaultMongoTypeMapper;
import org.springframework.data.mongodb.core.convert.MappingMongoConverter;
import org.springframework.data.mongodb.core.mapping.MongoMappingContext;
import org.springframework.util.StringUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import java.lang.reflect.Method;
import com.mongodb.ConnectionString;

import org.springframework.data.mongodb.MongoDatabaseFactory;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.SimpleMongoClientDatabaseFactory;
import org.springframework.data.mongodb.core.convert.MongoConverter;