1. 程式人生 > >SpringBoot+MongoDB實現多租戶切換

SpringBoot+MongoDB實現多租戶切換

目標:實現不同分組的使用者訪問自己的專用資料庫(資料庫在同一個伺服器).例:A公司使用者,僅可訪問A公司的專用資料庫

要求:MongoDB + SpringBoot

Pom依賴:

<parent>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-parent</artifactId>
    <version>2.0.3.RELEASE</version>
    <relativePath/> <!-- lookup parent from repository -->
</parent>

<properties>
    <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
    <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
    <java.version>1.8</java.version>
</properties>

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

    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-test</artifactId>
        <scope>test</scope>
    </dependency>
    <dependency>
        <groupId>org.springframework.data</groupId>
        <artifactId>spring-data-mongodb</artifactId>
    </dependency>
    <dependency>
        <groupId>org.mongodb</groupId>
        <artifactId>mongo-java-driver</artifactId>
    </dependency>
    <!--新增REST依賴-->
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-rest</artifactId>
    </dependency>
</dependencies>

1.自定義MongoDBFactoty

/**
 * 類的描述:
 * <p>自定義MongoDB Factory </p>
 * 自定義MongoDb 配置工廠,可實現多租戶切換訪問。
 *
 * @author jei0439
 * @date 2018/9/14 16:30
 */
public class MultiTenantMongoDbFactory extends SimpleMongoDbFactory {

    private static final Logger logger = LoggerFactory.getLogger(MultiTenantMongoDbFactory.class);

    /**
     * 預設資料庫名稱
     **/
    private final String defaultName;

    /**
     * MongoDB模板類
     **/
    private MongoTemplate mongoTemplate;

    /**
     * 使用者所線上程使用資料庫集合
     **/
    private static final ThreadLocal<String> dbName = new ThreadLocal<String>();
    /**
     *
     **/
    private static final HashMap<String, Object> databaseIndexMap = new HashMap<String, Object>();

    public MultiTenantMongoDbFactory(final MongoClient mongo, final String defaultDatabaseName) {
        super(mongo, defaultDatabaseName);
        logger.debug("Instantiating " + MultiTenantMongoDbFactory.class.getName() + " with default database name: " + defaultDatabaseName);
        this.defaultName = defaultDatabaseName;
    }

    /**
     *
     * 功能描述:  dirty but ... what can I do?
     *
     * @param
     * @author jei0439
     * @date 2018/9/17 10:55
     */
    public void setMongoTemplate(final MongoTemplate mongoTemplate) {
        Assert.isNull(this.mongoTemplate, "You can set MongoTemplate just once");
        this.mongoTemplate = mongoTemplate;
    }

    /**
     *
     * 功能描述: 將databaseName放入dbName
     *
     * @param databaseName database/scheme名稱
     * @author jei0439
     * @date 2018/9/17 10:56
     */
    public static void setDatabaseNameForCurrentThread(final String databaseName) {
        logger.debug("Switching to database: " + databaseName);
        dbName.set(databaseName);
    }

    /**
     *
     * 功能描述: 清空dbName
     *
     * @author jei0439
     * @date 2018/9/17 10:57
     */
    public static void clearDatabaseNameForCurrentThread() {
        if (logger.isDebugEnabled()) {
            logger.debug("Removing database [" + dbName.get() + "]");
        }
        dbName.remove();
    }

    @Override
    public MongoDatabase getDb() {
        final String tlName = dbName.get();
        final String dbToUse = (tlName != null ? tlName : this.defaultName);
        logger.debug("Acquiring database: " + dbToUse);
        createIndexIfNecessaryFor(dbToUse);
        return super.getDb(dbToUse);
    }

    /**
     *
     * 功能描述: TODO: 沒搞明白作用
     *
     * @param database  database/scheme 名稱
     * @author jei0439
     * @date 2018/9/17 10:58
     */
    private void createIndexIfNecessaryFor(final String database) {
        if (this.mongoTemplate == null) {
            logger.error("MongoTemplate is null, will not create any index.");
            return;
        }
//        sync and init once
        boolean needsToBeCreated = false;
        synchronized (MultiTenantMongoDbFactory.class) {
            final Object syncObj = databaseIndexMap.get(database);
            if (syncObj == null) {
                databaseIndexMap.put(database, new Object());
                needsToBeCreated = true;
            }
        }
//        make sure only one thread enters with needsToBeCreated = true
        synchronized (databaseIndexMap.get(database)) {
            if (needsToBeCreated) {
                logger.debug("Creating indices for database name=[" + database + "]");
                createIndexes();
                logger.debug("Done with creating indices for database name=[" + database + "]");
            }
        }
    }

    private void createIndexes() {
        final MongoMappingContext mappingContext = (MongoMappingContext) this.mongoTemplate.getConverter().getMappingContext();
        final MongoPersistentEntityIndexResolver indexResolver = new MongoPersistentEntityIndexResolver(mappingContext);
        for (BasicMongoPersistentEntity<?> persistentEntity : mappingContext.getPersistentEntities()) {
            checkForAndCreateIndexes(indexResolver, persistentEntity);
        }
    }

    private void checkForAndCreateIndexes(final MongoPersistentEntityIndexResolver indexResolver, final MongoPersistentEntity<?> entity) {
//        make sure its a root document
        if (entity.findAnnotation(Document.class) != null) {
            for (IndexDefinitionHolder indexDefinitionHolder : indexResolver.resolveIndexFor(entity.getTypeInformation())) {
//                work because of javas reentered lock feature
                this.mongoTemplate.indexOps(entity.getType()).ensureIndex(indexDefinitionHolder);
            }
        }
    }
}

2.MongoConfig

/**
 * 類的描述:
 * <p>自定義MongoDB Config </p>
 * Spring 此處檔案配置初始化 MongoDB 配置、自定義MongoDB Factory
 *
 * @author jei0439
 * @date 2018/9/14 16:30
 */
@Configuration
@EnableAutoConfiguration
public class MongoConfig {
    private final MongoClientOptions options;
    private final MongoClientFactory factory;
    private MongoClient mongo;

    /**
     *
     * 功能描述: 建構函式  直接使用MongoAutoConfiguration建構函式 引數初始化由Spring完成
     *
     * @param properties MongoDB 配置資料
     * @param options  TODO: 待完善 斷點跟蹤Spring框架傳參進來是null
     * @param environment TODO: 待完善 斷點跟蹤Spring框架傳參進來是null
     * @author jei0439
     * @date 2018/9/17 10:48
     */
    public MongoConfig(MongoProperties properties, ObjectProvider<MongoClientOptions> options, Environment environment) {
        this.options = (MongoClientOptions)options.getIfAvailable();
        this.factory = new MongoClientFactory(properties, environment);
    }
    /**
     *
     * 功能描述: 覆蓋預設的MongoDbFactory
     *
     * @author jei0439
     * @date 2018/9/17 10:51
     */
    @Bean
    MultiTenantMongoDbFactory mongoDbFactory() {
        this.mongo = this.factory.createMongoClient(this.options);
        MultiTenantMongoDbFactory mongoDbFactory = new MultiTenantMongoDbFactory(mongo,"pvQC_dev");
        return mongoDbFactory;
    }

    /**
     *
     * 功能描述: Create MongoTemplate
     *
     * @author jei0439
     * @date 2018/9/17 10:51
     */
    @Bean
    public MongoTemplate mongoTemplate() throws Exception {
        return new MongoTemplate(mongoDbFactory());
    }

}

3.測試Entity

/**
 * 類的描述:
 * <p></p>
 *
 * @author jei0439
 * @date 2018/9/17 10:30
 */

public class User {
    @Id
    private Long id;

    private String name;

    public Long getId() {
        return id;
    }

    public void setId(Long id) {
        this.id = id;
    }

    public String getName() {
        return name;
    }

    public void setName(String name) {
        this.name = name;
    }
}

4.Repositoty

/**
 * 類的描述:
 * <p></p>
 *
 * @author jei0439
 * @date 2018/9/17 10:31
 */
@Component
public interface Repository extends MongoRepository<User,Long> {
    User getById(Long id);
    User getByName(String name);
}

5.Application (隨便測試)

@SpringBootApplication
public class DemoApplication  implements CommandLineRunner {



    @Autowired
    Repository repository;



    @Override
    public void run(final String... args) throws Exception {
//        add something to the default database (test)
        User user = new User();
        user.setName("張三");
        user.setId(1L);
//      this.repository.save(user);

        System.out.println("data from test: " + this.repository.findAll());
//        okay? fine. - lets switch the database
        MultiTenantMongoDbFactory.setDatabaseNameForCurrentThread("test666");

        this.repository.save(user);
//        should be empty
        System.out.println("data from test666: " + this.repository.findAll());

//        switch back and clean up
        MultiTenantMongoDbFactory.setDatabaseNameForCurrentThread("test");
        this.repository.save(user);
    }


    public static void main(String[] args) {
        SpringApplication.run(DemoApplication.class, args);
    }
}

執行測試結果