背景

业务侧分库分表
根据账户切换不同数据库，多个账户在同一个库。例如0001-0005的客户在一个库。
根据账户分表，账户名_表名
其实使用mongodb不用业务侧做这种分库分表，历史原因，还是实现了。

分库

springboot mongodb 的 MongoTemplate 默认是一个database。我们需要动态切换。所以需要改造。

通过跟踪MongoTemplate构造我们发现，有一个关键的类MongoDatabaseFactory database工厂。

这个接口定义了一个方法getMongoDatabase()。而且他有一个简单的实现SimpleMongoClientDatabaseFactory。其中getMongoDatabase()获取的就是我们配置文件中的数据库名.

// 默认的实现，getDefaultDatabaseName() 就是配置文件中的 database
public MongoDatabase getMongoDatabase() throws DataAccessException {
	return getMongoDatabase(getDefaultDatabaseName());
}

按照SimpleMongoClientDatabaseFactory自己实现一个Factory, 其他不变，重写getMongoDatabase()

/**
 * 动态的 DynamicMongoDbFactory
 *
 * @author fengyi
 * @version 1.0
 * @date 2021/3/13 18:41
 */
public class DynamicMongoDbFactory extends MongoDatabaseFactorySupport<MongoClient>
        implements DisposableBean {

    private final static Logger LOG = LoggerFactory.getLogger(DynamicMongoDbFactory.class);

    /**
     * Creates a new {@link DynamicMongoDbFactory} instance for the given {@code connectionString}.
     *
     * @param connectionString connection coordinates for a database connection. Must contain a database name and must not
     *                         be {@literal null} or empty.
     * @see <a href="https://docs.mongodb.com/manual/reference/connection-string/">MongoDB Connection String reference</a>
     */
    public DynamicMongoDbFactory(String connectionString) {
        this(new ConnectionString(connectionString));
    }

    /**
     * Creates a new {@link DynamicMongoDbFactory} instance from the given {@link MongoClient}.
     *
     * @param connectionString connection coordinates for a database connection. Must contain also a database name and not
     *                         be {@literal null}.
     */
    public DynamicMongoDbFactory(ConnectionString connectionString) {
        this(MongoClients.create(connectionString), connectionString.getDatabase(), true);
    }

    /**
     * Creates a new {@link DynamicMongoDbFactory} instance from the given {@link MongoClient}.
     *
     * @param mongoClient  must not be {@literal null}.
     * @param databaseName must not be {@literal null} or empty.
     */
    public DynamicMongoDbFactory(MongoClient mongoClient, String databaseName) {
        this(mongoClient, databaseName, false);
    }

    /**
     * Creates a new {@link DynamicMongoDbFactory} instance from the given {@link MongoClient}.
     *
     * @param mongoClient          must not be {@literal null}.
     * @param databaseName         must not be {@literal null} or empty.
     * @param mongoInstanceCreated
     */
    DynamicMongoDbFactory(MongoClient mongoClient, String databaseName, boolean mongoInstanceCreated) {
        super(mongoClient, databaseName, mongoInstanceCreated, new MongoExceptionTranslator());
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.MongoDbFactory#getSession(com.mongodb.ClientSessionOptions)
     */
    @Override
    public ClientSession getSession(ClientSessionOptions options) {
        return getMongoClient().startSession(options);
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.core.MongoDbFactoryBase#closeClient()
     */
    @Override
    protected void closeClient() {
        getMongoClient().close();
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.mongodb.core.MongoDbFactoryBase#doGetMongoDatabase(java.lang.String)
     */
    @Override
    protected MongoDatabase doGetMongoDatabase(String dbName) {
        return getMongoClient().getDatabase(dbName);
    }

    @Override
    public MongoDatabase getMongoDatabase() throws DataAccessException {
        String db = DbHolder.get() == null ? getDefaultDatabaseName() : DbHolder.get();
        if (LOG.isDebugEnabled()) {
            LOG.debug("DB Holder: {}", db);
        }
        return super.getMongoDatabase(db);
    }
}

分库信息上下文 ThreadLocal

/**
 * 分库信息上下文
 *
 * @author fengyi
 * @version 1.0
 * @date 2021/3/15 13:55
 */
public class DbHolder {
    private static final ThreadLocal<String> DB_NAME = new InheritableThreadLocal<>();

    public static void set(String dbName){
        DB_NAME.set(dbName);
    }

    public static String get(){
        return DB_NAME.get();
    }

    public static void remove(){
        DB_NAME.remove();
    }
}

注入bean

    /**
     * 动态分库
     * @param mongoClient
     * @return
     */
    @Primary
    @Bean
    public MongoDatabaseFactory dynamicMongoDbFactory(MongoProperties mongoProperties, MongoClient mongoClient){
        return new DynamicMongoDbFactory(mongoClient, mongoProperties.getDatabase());
    }

除此之外还需要注入MongoClient驱动的配置，否则会启动报错

@Configuration(
        proxyBeanMethods = false
)
@ConditionalOnClass({MongoClient.class})
public class MongoClientConfig {
    public MongoClientConfig() {
    }

    @Bean
    @ConditionalOnMissingBean(MongoClient.class)
    public MongoClient mongo(MongoProperties properties, Environment environment,
                             ObjectProvider<MongoClientSettingsBuilderCustomizer> builderCustomizers,
                             ObjectProvider<MongoClientSettings> settings) {
        return new MongoClientFactory(properties, environment,
                builderCustomizers.orderedStream().collect(Collectors.toList()))
                .createMongoClient(settings.getIfAvailable());
    }
}

分表

关键类BasicMongoPersistentEntity, 重写获取集合的方法。直接上代码 ，不废话了

重写 MongoMappingContext 的 typeInformation方法

/**
 * 重写MongoMappingContext
 *
 * @author fengyi
 * @version 1.0
 * @date 2021/3/15 14:11
 */
public class MyMongoMappingContext extends MongoMappingContext {


    @Override
    protected <T> DynamicMongoPersistentEntity<T> createPersistentEntity(TypeInformation<T> typeInformation) {
        return new DynamicMongoPersistentEntity<T>(typeInformation);
    }
}

DynamicMongoPersistentEntity 继承BasicMongoPersistentEntity，重写getCollection()方法

/**
 * 动态的 表名
 *
 * @author fengyi
 * @version 1.0
 * @date 2021/3/15 13:49
 */
public class DynamicMongoPersistentEntity<T> extends BasicMongoPersistentEntity<T> {

    private final static Logger LOG = LoggerFactory.getLogger(DynamicMongoPersistentEntity.class);

    private final static String UNDERLINE = "_";

    /**
     * Creates a new {@link BasicMongoPersistentEntity} with the given {@link TypeInformation}. Will default the
     * collection name to the entities simple type name.
     *
     * @param typeInformation must not be {@literal null}.
     */
    public DynamicMongoPersistentEntity(TypeInformation<T> typeInformation) {
        super(typeInformation);
    }

    @Override
    public String getCollection() {
        String collection = super.getCollection();
        collection = AccountHolder.get() == null ? collection : AccountHolder.get().concat(UNDERLINE).concat(collection);
        if (LOG.isDebugEnabled()) {
            LOG.debug("Collection Holder: {}", collection);
        }
        return collection;
    }

}

AccountHolder 账号信息 ThreadLocal

/**
 * 账号信息上下文
 *
 * @author fengyi
 * @version 1.0
 * @date 2021/3/15 13:55
 */
public class AccountHolder {
    private static final ThreadLocal<String> ACCOUNT = new InheritableThreadLocal<>();

    public static void set(String account){
        ACCOUNT.set(account);
    }

    public static String get(){
        return ACCOUNT.get();
    }

    public static void remove(){
        ACCOUNT.remove();
    }
}

注入bean

    /**
     * 动态分表
     * @param applicationContext
     * @param properties
     * @param conversions
     * @return
     * @throws ClassNotFoundException
     */
    @Primary
    @Bean
    MyMongoMappingContext mongoMappingContext(ApplicationContext applicationContext, MongoProperties properties, MongoCustomConversions conversions) throws ClassNotFoundException {

        MyMongoMappingContext context = new MyMongoMappingContext();
        context.setInitialEntitySet((new EntityScanner(applicationContext)).scan(new Class[]{Document.class, Persistent.class}));
        Class<?> strategyClass = properties.getFieldNamingStrategy();
        if (strategyClass != null) {
            context.setFieldNamingStrategy((FieldNamingStrategy) BeanUtils.instantiateClass(strategyClass));
        }

        context.setSimpleTypeHolder(conversions.getSimpleTypeHolder());
        return context;
    }

测试结果

结束语

文章只是描述了核心原理，具体设置账户和分库的逻辑，要根据实际业务场景来实现。比如从http header中获取，设置到AccountHolder。而分库信息DbHolder的设置，则应该有一张路由表，用来查找这个账户在那个库。有问题可以留言。