com.querydsl.mongodb.document.AbstractFetchableMongodbQuery Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2020 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.querydsl.mongodb.document;

import com.mongodb.ReadPreference;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoCursor;
import com.mysema.commons.lang.CloseableIterator;
import com.querydsl.core.*;
import com.querydsl.core.types.Expression;
import com.querydsl.core.types.OrderSpecifier;
import com.querydsl.core.types.Path;
import com.querydsl.core.types.Predicate;
import org.bson.Document;
import org.jetbrains.annotations.Nullable;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.function.Function;

/**
 * {@link Fetchable} Mongodb query with a pluggable Document to Bean transformation.
 *
 * @param  result type
 * @param  concrete subtype
 * @author Mark Paluch
 */
public abstract class AbstractFetchableMongodbQuery>
        extends AbstractMongodbQuery implements Fetchable {

    private final Function transformer;

    private final MongoCollection collection;

    /**
     * Create a new MongodbQuery instance
     * @param collection
     * @param transformer result transformer
     * @param serializer serializer
     */
    public AbstractFetchableMongodbQuery(MongoCollection collection, Function transformer, MongodbDocumentSerializer serializer) {
        super(serializer);
        this.transformer = transformer;
        this.collection = collection;
    }

    /**
     * Iterate with the specific fields
     *
     * @param paths fields to return
     * @return iterator
     */
    public CloseableIterator iterate(Path... paths) {
        getQueryMixin().setProjection(paths);
        return iterate();
    }

    @Override
    public CloseableIterator iterate() {
        FindIterable cursor = createCursor();
        final MongoCursor iterator = cursor.iterator();

        return new CloseableIterator() {
            @Override
            public boolean hasNext() {
                return iterator.hasNext();
            }

            @Override
            public K next() {
                return transformer.apply(iterator.next());
            }

            @Override
            public void remove() {

            }

            @Override
            public void close() {
                iterator.close();
            }
        };
    }

    /**
     * Fetch with the specific fields
     *
     * @param paths fields to return
     * @return results
     */
    public List fetch(Path... paths) {
        getQueryMixin().setProjection(paths);
        return fetch();
    }

    @Override
    public List fetch() {
        try {
            FindIterable cursor = createCursor();
            List results = new ArrayList();
            for (Document document : cursor) {
                results.add(transformer.apply(document));
            }
            return results;
        } catch (NoResults ex) {
            return Collections.emptyList();
        }
    }

    /**
     * Fetch first with the specific fields
     *
     * @param paths fields to return
     * @return first result
     */
    public K fetchFirst(Path... paths) {
        getQueryMixin().setProjection(paths);
        return fetchFirst();
    }

    @Override
    public K fetchFirst() {
        try {
            FindIterable c = createCursor().limit(1);
            MongoCursor iterator = c.iterator();
            try {

                if (iterator.hasNext()) {
                    return transformer.apply(iterator.next());
                } else {
                    return null;
                }
            } finally {
                iterator.close();
            }
        } catch (NoResults ex) {
            return null;
        }
    }

    /**
     * Fetch one with the specific fields
     *
     * @param paths fields to return
     * @return first result
     */
    public K fetchOne(Path... paths) {
        getQueryMixin().setProjection(paths);
        return fetchOne();
    }

    @Override
    public K fetchOne() {
        try {
            Long limit = getQueryMixin().getMetadata().getModifiers().getLimit();
            if (limit == null) {
                limit = 2L;
            }

            FindIterable c = createCursor().limit(limit.intValue());
            MongoCursor iterator = c.iterator();
            try {

                if (iterator.hasNext()) {
                K rv = transformer.apply(iterator.next());
                if (iterator.hasNext()) {
                    throw new NonUniqueResultException();
                }
                return rv;
            } else {
                return null;
            }
            } finally {
                iterator.close();
            }


        } catch (NoResults ex) {
            return null;
        }
    }

    /**
     * Fetch results with the specific fields
     *
     * @param paths fields to return
     * @return results
     */
    public QueryResults fetchResults(Path... paths) {
        getQueryMixin().setProjection(paths);
        return fetchResults();
    }

    @Override
    public QueryResults fetchResults() {
        try {
            long total = fetchCount();
            if (total > 0L) {
                return new QueryResults(fetch(), getQueryMixin().getMetadata().getModifiers(), total);
            } else {
                return QueryResults.emptyResults();
            }
        } catch (NoResults ex) {
            return QueryResults.emptyResults();
        }
    }

    @Override
    public long fetchCount() {
        try {
            Predicate filter = createFilter(getQueryMixin().getMetadata());
            return collection.count(createQuery(filter));
        } catch (NoResults ex) {
            return 0L;
        }
    }

    protected FindIterable createCursor() {
        QueryMetadata metadata = getQueryMixin().getMetadata();
        Predicate filter = createFilter(metadata);
        return createCursor(collection, filter, metadata.getProjection(), metadata.getModifiers(), metadata.getOrderBy());
    }

    protected FindIterable createCursor(MongoCollection collection, @Nullable Predicate where,
                                                  Expression projection, QueryModifiers modifiers, List> orderBy) {

        ReadPreference readPreference = getReadPreference();
        MongoCollection collectionToUse = readPreference != null ? collection
                .withReadPreference(readPreference) : collection;
        FindIterable cursor = collectionToUse.find(createQuery(where))
                .projection(createProjection(projection));
        Integer limit = modifiers.getLimitAsInteger();
        Integer offset = modifiers.getOffsetAsInteger();

        if (limit != null) {
            cursor = cursor.limit(limit);
        }
        if (offset != null) {
            cursor = cursor.skip(offset);
        }
        if (orderBy.size() > 0) {
            cursor = cursor.sort(getSerializer().toSort(orderBy));
        }
        return cursor;
    }

    protected abstract MongoCollection getCollection(Class type);

    @Override
    protected List