Bladeren bron

新增 地理位置索引自动创建

zhaiyifei 10 maanden geleden
bovenliggende
commit
5b4ddf82bb
1 gewijzigde bestanden met toevoegingen van 37 en 0 verwijderingen
  1. 37 0
      src/main/java/com/yunfeiyun/agmp/iot/common/service/MongoService.java

+ 37 - 0
src/main/java/com/yunfeiyun/agmp/iot/common/service/MongoService.java

@@ -26,6 +26,7 @@ import org.springframework.data.domain.Sort;
 import org.springframework.data.mongodb.core.BulkOperations;
 import org.springframework.data.mongodb.core.MongoTemplate;
 import org.springframework.data.mongodb.core.aggregation.*;
+import org.springframework.data.mongodb.core.index.GeospatialIndex;
 import org.springframework.data.mongodb.core.query.Criteria;
 import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
@@ -35,6 +36,8 @@ import org.springframework.stereotype.Service;
 import java.util.*;
 import java.util.regex.Pattern;
 
+import static org.springframework.data.mongodb.core.index.GeoSpatialIndexType.GEO_2DSPHERE;
+
 @Data
 class AggregateTotalBean{
     private int total;
@@ -681,6 +684,40 @@ public class  MongoService<T>{
         }
     }
 
+    public void createGeoIndexes(Class<T> entity, List<IndexModel> indexModelList) {
+        TableName annotation = entity.getAnnotation(TableName.class);
+        String tableName = annotation.value();
+        if(!mongoTemplate.collectionExists(tableName)){
+            mongoTemplate.createCollection(tableName);
+        }
+        MongoCollection<Document> collection = mongoTemplate.getCollection(tableName);
+        ListIndexesIterable<Document> indexInfoList = collection.listIndexes();
+        Set<String> indexNameSet = new HashSet<>();
+        for (Document indexInfo : indexInfoList) {
+            indexNameSet.add(indexInfo.getString("name"));
+        }
+        log.info("{} 已存在索引: {}", tableName, indexNameSet);
+        List<String> geoIndexNameList = new ArrayList<>();
+        for(IndexModel indexModel : indexModelList) {
+            Bson keys = indexModel.getKeys();
+            StringBuilder indexName = new StringBuilder();
+            BsonDocument bsonDocument = keys.toBsonDocument();
+            for (Map.Entry<String, BsonValue> entry : bsonDocument.entrySet()) {
+                String key = entry.getKey();
+                indexName.append(key).append("_").append("2dsphere");
+            }
+            String indexNameStr = String.valueOf(indexName);
+            if(!indexNameSet.contains(indexNameStr)) {
+                geoIndexNameList.add(indexNameStr);
+                log.info("创建mongodb {} 索引: {}", tableName, indexNameStr);
+            }
+        }
+        for(String indexName : geoIndexNameList){
+            GeospatialIndex geoIndex = new GeospatialIndex(indexName).typed(GEO_2DSPHERE);
+            mongoTemplate.indexOps(tableName).ensureIndex(geoIndex);
+        }
+    }
+
     public void updateBatch(Class<T> entity, List<MgUpdateBatchDto> updateBatchDtoList) {
         TableName annotation = entity.getAnnotation(TableName.class);
         if (annotation == null || StringUtils.isEmpty(annotation.value())) {