mongodb动态分表

//测试类
@RestController
@RequestMapping("/test")
@Slf4j
public class TestController {
    @Autowired(required = false)
    SelfBeanToSpring selfBeanToSpring;
    @Autowired
    RepositoryPlusUtil repositoryPlusUtil;
    @GetMapping("/spel")
    @ResponseBody
    public Object distinct() {
        List<Esl> lh001 = repositoryPlusUtil.getRepositorySpel(EslRepositoryPlus.class, "lh001").findByEslCodeIn(Arrays.asList(1L));
        return lh001;
    }
}


//工具类
@Component
@Slf4j
public class RepositoryPlusUtil {
    @Autowired
    MongoTemplate mongoTemplate;
    @Autowired
    MongoSpelDocumentNameUtil mongoSpelDocumentNameUtil;
    /**
     * 存放生成的分表MongoRepository
     */
    private Map<Class, MongoRepository> map = new ConcurrentHashMap<>();
    /**
     * map中有分表名称则不需要执行创建索引
     */
    private Map<String, Integer> mapIndexCreateFlagMap = new ConcurrentHashMap<>();

    /**
     * 获取分表名称并进行索引创建
     *
     * @param prefix 分表前缀
     * @param clazz  mongo对应数据库的实体类,类上有@Document(collection = "collectionName")
     * @return
     */
    public <T> String getCustomCollectionName(String prefix, Class<T> clazz) {
        //分表名
        String customCollectionName = prefix + "_" + MongoUtil.getCollecitonName(clazz);
        if (mapIndexCreateFlagMap.get(customCollectionName) == null) {
            //创建索引
            createIndex(clazz, customCollectionName);
        }
        return customCollectionName;
    }

    /**
     * 利用spel支持动态设置document的name,从而实现分库分表
     *
     * @param clazz
     * @param prefix
     * @param <T>
     * @return
     */
    public <T> T getRepositorySpel(Class<T> clazz, String prefix) {
        try {
            Type[] genericInterfaces = clazz.getGenericInterfaces();
            Type genericInterface = genericInterfaces[0];
            String documentEntityName = genericInterface.getTypeName().split("<")[1].split(",")[0];
            org.springframework.data.mongodb.core.mapping.Document document = (org.springframework.data.mongodb.core.mapping.Document) Class.forName(documentEntityName).getAnnotations()[0];
            /*
            @Document(collection = "#{@mongoSpelDocumentNameUtil.get()}_esl_info")
            线程隔离
            @Component("mongoSpelDocumentNameUtil")
            public class MongoSpelDocumentNameUtil {
                private final  ThreadLocal<String> holder= new ThreadLocal<>();
                public String get(){
                    String value = holder.get();
                    return ObjectUtils.isEmpty(value)?"":value;
                }
                public void set(String prefix){
                    if(prefix==null){
                        throw new RuntimeException("tenantId not null.");
                    }
                    holder.set(prefix);
                }
                public void clear(){
                    holder.remove();
                }
            }
             */
            String name = document.collection().replace("#{@mongoSpelDocumentNameUtil.get()}_", "");
            log.info(name);
            //分表名
            String customCollectionName = prefix + "_" + name;
            if (mapIndexCreateFlagMap.get(customCollectionName) == null) {
                //创建索引
                createIndex(clazz, customCollectionName);
            }
        } catch (Exception e) {
            log.error("", e);
        }
        mongoSpelDocumentNameUtil.set(prefix + "_");
        return SpringContextUtils.getApplicationContext().getBean(clazz);
    }


    /**
     * @param clazz  mongo对应数据库的实体类,类上有@Document(collection = "collectionName")
     * @param prefix 分表前缀
     * @return
     */
    public <T> MongoRepository<T, String> getRepositoryPlus(Class<T> clazz, String prefix) {
        log.info("className={}", clazz.getName());
        try {
            //分表名
            String customCollectionName = prefix + "_" + MongoUtil.getCollecitonName(clazz);
            if (mapIndexCreateFlagMap.get(customCollectionName) == null) {
                //创建索引
                createIndex(clazz, customCollectionName);
            }
            //生成分表对应的repository
            MongoRepository mongoRepository = map.get(clazz);
            if (mongoRepository == null) {
                synchronized (RepositoryPlusUtil.class) {
                    mongoRepository = map.get(clazz);
                    if (mongoRepository == null) {
                        MongoPersistentEntity<T> persistentEntity = (MongoPersistentEntity<T>) mongoTemplate.getConverter().getMappingContext().getPersistentEntity(clazz);
                        MongoEntityInformation<T, String> mongoEntityInformation = new MappingMongoEntityInformation<>(persistentEntity, customCollectionName);
                        SimpleMongoRepository<T, String> sysConfigStringSimpleMongoRepository = new SimpleMongoRepository<>(mongoEntityInformation, mongoTemplate);
                        map.put(clazz, sysConfigStringSimpleMongoRepository);
                        return sysConfigStringSimpleMongoRepository;
                    }
                }
            }
            return mongoRepository;
        } catch (Exception e) {
            log.error("", e);
            return null;
        }
    }

    /**
     * 创建索引
     *
     * @param clazz
     * @param customCollectionName
     * @param <T>
     */
    private <T> void createIndex(Class<T> clazz, String customCollectionName) {
        //获取分表所有的index
        IndexOperations indexOperations = mongoTemplate.indexOps(customCollectionName);
        List<IndexInfo> indexInfos = indexOperations.getIndexInfo();
        List<String> indexNames = new ArrayList<>();
        if (!CollectionUtil.isEmpty(indexInfos)) {
            indexInfos.forEach(indexInfo -> indexNames.add(indexInfo.getName()));
        }
        //单索引 key=field value=Indexed
        Map<String, Indexed> indexedMap = new HashMap<>();
        //联合索引
        List<CompoundIndex> compoundIndexList = new ArrayList<>();
        //本类获取索引
        extracted(clazz, indexedMap, compoundIndexList);
        //父类获取索引
        Class<? super T> superclass = clazz.getSuperclass();
        if (superclass != null && !ObjectUtils.isEmpty(superclass.getPackage().getName())) {
            extracted(superclass, indexedMap, compoundIndexList);
        }
        //单个索引创建
        indexedMap.entrySet().forEach(indexed -> {
            String indexName = ObjectUtils.isEmpty(indexed.getValue().name()) ? indexed.getKey() : indexed.getValue().name();
            String indexOnField = indexed.getKey();
            if (!indexNames.contains(indexName)) {
                Index index = new Index();
                switch (indexed.getValue().direction().name()) {
                    case "ASCENDING":
                        index.on(indexOnField, Sort.Direction.ASC);
                        break;
                    case "DESCENDING":
                        index.on(indexOnField, Sort.Direction.DESC);
                        break;
                    default:
                        index.on(indexOnField, Sort.Direction.ASC);
                }
                index.expire(indexed.getValue().expireAfterSeconds(), TimeUnit.SECONDS);
                index.named(indexName);
                log.info("create index={},on {}", indexName, customCollectionName);
                indexOperations.ensureIndex(index);
            }
        });
        //复合索引创建
        compoundIndexList.forEach(indexed -> {
            try {
                StringJoiner stringJoiner = new StringJoiner("_");
                String indexName = ObjectUtils.isEmpty(indexed.name()) ? stringJoiner.toString() : indexed.name();
                if (!indexNames.contains(indexName)) {
                    Document indexDocument = JacksonMapper.getInstance().readValue(indexed.def().replaceAll("'", "\""), Document.class);
                    indexDocument.keySet().forEach(key -> stringJoiner.add(key).add(indexDocument.get(key).toString()));
                    CompoundIndexDefinition compoundIndexDefinition = new CompoundIndexDefinition(indexDocument);
                    compoundIndexDefinition.named(indexName);
                    indexOperations.ensureIndex(compoundIndexDefinition);
                    log.info("create compoundIndex={},on {}", indexName, customCollectionName);
                }
            } catch (Exception e) {
                log.error("", e);
            }
        });
        mapIndexCreateFlagMap.put(customCollectionName, 1);
    }

    /**
     * 索引获取
     *
     * @param clazz
     * @param indexedMap
     * @param compoundIndexList
     * @param <T>
     */
    private <T> void extracted(Class<T> clazz, Map<String, Indexed> indexedMap, List<CompoundIndex> compoundIndexList) {
        //单索引获取
        Field[] declaredFields = clazz.getDeclaredFields();
        if (declaredFields != null) {
            for (int i = 0; i < declaredFields.length; i++) {
                Annotation[] annotations = declaredFields[i].getAnnotations();
                if (annotations != null) {
                    for (int j = 0; j < annotations.length; j++) {
                        Annotation annotation = annotations[j];
                        if (annotation instanceof Indexed) {
                            Indexed indexed = (Indexed) annotation;
                            indexedMap.put(declaredFields[i].getName(), indexed);
                        }
                    }
                }
            }
        }
        //联合索引获取
        Annotation[] classAnnotations = clazz.getAnnotations();
        if (classAnnotations != null) {
            for (int i = 0; i < classAnnotations.length; i++) {
                Annotation annotation = classAnnotations[i];
                if (annotation instanceof CompoundIndex) {
                    //单个联合索引
                    CompoundIndex compoundIndex = (CompoundIndex) annotation;
                    compoundIndexList.add(compoundIndex);
                }
                if (annotation instanceof CompoundIndexes) {
                    //多个联合索引
                    CompoundIndexes compoundIndexes = (CompoundIndexes) annotation;
                    CompoundIndex[] value = compoundIndexes.value();
                    if (value != null && value.length > 0) {
                        for (int i1 = 0; i1 < value.length; i1++) {
                            compoundIndexList.add(value[i1]);
                        }
                    }
                }
            }
        }
    }

    public static void main(String[] args) throws JsonProcessingException {
        String str = "{'scGroup': 1, 'scKey': 1}".replaceAll("'", "\"");
        Map s = JacksonMapper.getInstance().readValue(str, Map.class);
        System.out.println(s);
    }
}


//Repository
@Repository
@Scope("prototype")
public interface EslRepositoryPlus extends ResourceRepository<EslSpel, String> {
    @Query(fields = "{'eslCode':1,'storeCode':1,'status':1}")
    List<Esl> findByEslCodeIn(List<Long> codes);
}


//Repository对应的数据库实体类
@Getter
@Setter
@Document(collection = "#{@mongoSpelDocumentNameUtil.get()}_esl")
public class EslSpel extends CommonParam{
    private Long eslCode;
}