MongoTemplate使用案例

package cn.idea360.assistant.dev.mongo;

import com.mongodb.ConnectionString;
import org.bson.Document;
import org.bson.types.ObjectId;
import org.junit.Before;
import org.junit.Test;
import org.springframework.data.domain.Sort;
import org.springframework.data.mongodb.core.MongoTemplate;
import org.springframework.data.mongodb.core.SimpleMongoClientDatabaseFactory;
import org.springframework.data.mongodb.core.aggregation.Aggregation;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;
import org.springframework.data.util.CloseableIterator;
import org.springframework.util.CollectionUtils;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author cuishiying
 * @date 2024-11-06
 */
public class MongoTest {

    private MongoTemplate mongoTemplate;

    @Before
    public void newMongoTemplate() {
        ConnectionString connectionString = new ConnectionString("mongodb://coreuser:uil5sQpS9ODA@10.10.10.206:27017/coredb");
        SimpleMongoClientDatabaseFactory mongoClientDatabaseFactory = new SimpleMongoClientDatabaseFactory(connectionString);
        mongoTemplate = new MongoTemplate(mongoClientDatabaseFactory);
    }

    /**
     * 使用 MongoTemplate 分页查询.
     *
     * <p>
     * 在 MongoDB 中使用 skip() 处理大量数据时会有性能问题。因为 MongoDB 需要先扫描跳过的文档，再返回所需的文档，随着 skip 值的增加，查询会变得越来越慢。
     * </p>
     */
    @Test
    public void pageQuery() {
        int pageSize = 100;
        long skip = 0;
        String collectionName = "monthly_bills";

        Criteria criteria = new Criteria();
        criteria.and("billingMonth").is(202410);
        Query query = Query.query(criteria);

        query.with(Sort.by(Sort.Direction.DESC, "totalCalls"));
        query.skip(skip).limit(pageSize);

        List<Document> rows = mongoTemplate.find(query, Document.class, collectionName);
        System.out.println("数据量: " + rows.size());
    }

    /**
     * 使用游标查询.
     *
     * <p>内存占用低，适合大数据量查询</p>
     */
    @Test
    public void queryByCursor() {
        int batchSize = 100;
        String collectionName = "monthly_bills";

        Criteria criteria = new Criteria();
        criteria.and("billingMonth").is(202410);
        Query query = Query.query(criteria);
        query.with(Sort.by(Sort.Direction.DESC, "totalCalls"));

        long totalCompanies = mongoTemplate.count(query, Document.class, collectionName);
        System.out.println("totalCompanies: " + totalCompanies);

        AtomicInteger processedCount = new AtomicInteger(0);
        try (CloseableIterator<Document> iterator = mongoTemplate.stream(
                query, Document.class, collectionName)) {
            while (iterator.hasNext()) {
                Document document = iterator.next();
                int current = processedCount.incrementAndGet();
                if (current % batchSize == 0) {
                    double progress = (double) current / totalCompanies * 100;
                    System.out.printf("统计进度: %s/%s (%s)%n",
                            current, totalCompanies, String.format("%.2f", progress));
                }
            }
        }

        System.out.println("数据量: " + processedCount.get());
    }

    /**
     * 使用聚合查询.
     *
     * <p>支持复杂的聚合操作, 可以进行分组、计算、转换等操作, 性能取决于聚合管道的复杂度</p>
     */
    @Test
    public void queryByAggregation() {
        int batchSize = 100;
        String collectionName = "monthly_bills";

        Criteria criteria = new Criteria();
        criteria.and("billingMonth").is(202410);

        long totalCompanies = mongoTemplate.aggregate(
                Aggregation.newAggregation(
                        Aggregation.match(criteria),
                        Aggregation.group("cid")
                ),
                collectionName,
                Document.class
        ).getMappedResults().size();
        System.out.println("totalCompanies: " + totalCompanies);

        Aggregation aggregation = Aggregation.newAggregation(
                // 可选：匹配条件（如时间范围）
                Aggregation.match(criteria),
                // 按公司ID分组并计数
                Aggregation.group("cid")
                        .count().as("count")
                        .sum("amount").as("totalAmount")
                        .avg("amount").as("avgAmount")
                        .min("createTime").as("firstTime")
                        .max("createTime").as("lastTime"),

                // 可选：按数量排序
                Aggregation.sort(Sort.Direction.DESC, "count")
        );

        AtomicInteger processedCount = new AtomicInteger(0);
        try (CloseableIterator<Document> iterator = mongoTemplate.aggregateStream(
                aggregation, collectionName, Document.class)) {

            while (iterator.hasNext()) {
                Document document = iterator.next();
                int current = processedCount.incrementAndGet();
                if (current % batchSize == 0) {
                    double progress = (double) current / totalCompanies * 100;
                    System.out.printf("统计进度: %s/%s (%s)%n",
                            current, totalCompanies, String.format("%.2f", progress));
                }
            }

        }

        System.out.println("数据量: " + processedCount.get());
    }

    /**
     * 多条件复杂查询
     */
    @Test
    public void complexQuery() {
        String collectionName = "monthly_bills";

        // 复杂条件
        List<Criteria> criteriaList = new ArrayList<>();
        criteriaList.add(Criteria.where("billingMonth").is(202410));
        Criteria buildCriteria = new Criteria().andOperator(criteriaList.toArray(new Criteria[0]));

        List<Sort.Order> orders = new ArrayList<>();
        // 始终添加_id排序保证分页稳定性
        orders.add(Sort.Order.asc("_id"));
        Sort buildSort = Sort.by(orders);

        Query query = Query.query(buildCriteria).with(buildSort);
        long count = mongoTemplate.count(query, Document.class, collectionName);
        System.out.println("数据量: " + count);
    }
}