org.opencb.commons.datastore.mongodb.MongoDBNativeQuery.java Source code

Java tutorial

Introduction

Here is the source code for org.opencb.commons.datastore.mongodb.MongoDBNativeQuery.java

Source

/*
 * Copyright 2015-2017 OpenCB
 *
 * 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 org.opencb.commons.datastore.mongodb;

import com.mongodb.WriteConcern;
import com.mongodb.bulk.BulkWriteResult;
import com.mongodb.client.AggregateIterable;
import com.mongodb.client.DistinctIterable;
import com.mongodb.client.FindIterable;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.model.*;
import com.mongodb.client.result.DeleteResult;
import com.mongodb.client.result.UpdateResult;
import org.bson.Document;
import org.bson.conversions.Bson;
import org.opencb.commons.datastore.core.QueryOptions;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.TimeUnit;

import static org.opencb.commons.datastore.mongodb.MongoDBQueryUtils.getProjection;
import static org.opencb.commons.datastore.mongodb.MongoDBQueryUtils.parseQueryOptions;

/**
 * Created by imedina on 28/03/14.
 */
public class MongoDBNativeQuery {

    private final MongoCollection<Document> dbCollection;

    MongoDBNativeQuery(MongoCollection<Document> dbCollection) {
        this.dbCollection = dbCollection;
    }

    MongoCollection<Document> getDbCollection() {
        return dbCollection;
    }

    public long count() {
        return dbCollection.count();
    }

    public long count(Bson query) {
        //        CountOptions c = new CountOptions().
        return dbCollection.count(query);
    }

    public DistinctIterable<Document> distinct(String key) {
        return distinct(key, null, Document.class);
    }

    public <T> DistinctIterable distinct(String key, Class<T> resultClass) {
        return distinct(key, null, resultClass);
    }

    public <T> DistinctIterable<T> distinct(String key, Bson query, Class<T> resultClass) {
        return dbCollection.distinct(key, query, resultClass);
    }

    public FindIterable<Document> find(Bson query, QueryOptions options) {
        return find(query, null, options);
    }

    public AggregateIterable<Document> aggregate(List<? extends Bson> operations) {
        return aggregate(operations, null);
    }

    public AggregateIterable<Document> aggregate(List<? extends Bson> operations, QueryOptions options) {
        // we need to be sure that the List is mutable
        List<Bson> bsonOperations = new ArrayList<>(operations);
        parseQueryOptions(bsonOperations, options);
        return (bsonOperations.size() > 0) ? dbCollection.aggregate(bsonOperations) : null;
    }

    public FindIterable<Document> find(Bson query, Bson projection, QueryOptions options) {

        if (projection == null) {
            projection = getProjection(projection, options);
        }

        FindIterable<Document> findIterable = dbCollection.find(query).projection(projection);

        int limit = (options != null) ? options.getInt(QueryOptions.LIMIT, 0) : 0;
        if (limit > 0) {
            findIterable.limit(limit);
        }

        int skip = (options != null) ? options.getInt(QueryOptions.SKIP, 0) : 0;
        if (skip > 0) {
            findIterable.skip(skip);
        }

        Object sortObject = (options != null) ? options.get(QueryOptions.SORT) : null;
        if (sortObject != null) {
            if (sortObject instanceof Bson) {
                findIterable.sort(((Bson) sortObject));
            } else if (sortObject instanceof String) {
                String order = options.getString(QueryOptions.ORDER, "DESC");
                if (order.equalsIgnoreCase(QueryOptions.ASCENDING) || order.equalsIgnoreCase("ASC")
                        || order.equals("1")) {
                    findIterable.sort(Sorts.ascending(((String) sortObject)));
                } else {
                    findIterable.sort(Sorts.descending(((String) sortObject)));
                }
            }
        }

        if (options != null && options.containsKey(MongoDBCollection.BATCH_SIZE)) {
            findIterable.batchSize(options.getInt(MongoDBCollection.BATCH_SIZE, 20));
        }

        if (options != null && options.containsKey(QueryOptions.TIMEOUT)) {
            findIterable.maxTime(options.getLong(QueryOptions.TIMEOUT), TimeUnit.MILLISECONDS);
        }

        return findIterable;
    }

    /**
     * This method insert a single document into a collection. Params w and wtimeout are read from QueryOptions.
     *
     * @param document The new document to be inserted
     * @param options  Some options like timeout
     */
    public void insert(Document document, QueryOptions options) {
        int writeConcern = 1;
        int ms = 0;
        if (options != null && (options.containsKey("w") || options.containsKey("wtimeout"))) {
            // Some info about params: http://api.mongodb.org/java/current/com/mongodb/WriteConcern.html
            //            return dbCollection.insert(dbObject, new WriteConcern(options.getInt("w", 1), options.getInt("wtimeout", 0)));
            writeConcern = options.getInt("w", 1);
            ms = options.getInt("wtimeout", 0);
        }
        dbCollection.withWriteConcern(new WriteConcern(writeConcern, ms));
        dbCollection.insertOne(document);
    }

    /**
     * This method insert a list of documents into a collection. Params w and wtimeout are read from QueryOptions.
     *
     * @param documentList The new list of documents to be inserted
     * @param options      Some options like timeout
     * @return A BulkWriteResult from MongoDB API
     */
    public BulkWriteResult insert(List<Document> documentList, QueryOptions options) {
        List<WriteModel<Document>> actions = new ArrayList<>(documentList.size());
        for (Document document : documentList) {
            actions.add(new InsertOneModel<>(document));
        }

        int writeConcern = 1;
        int ms = 0;
        if (options != null && (options.containsKey("w") || options.containsKey("wtimeout"))) {
            writeConcern = options.getInt("w", 1);
            ms = options.getInt("wtimeout", 0);
        }
        dbCollection.withWriteConcern(new WriteConcern(writeConcern, ms));

        return dbCollection.bulkWrite(actions, new BulkWriteOptions().ordered(false));
    }

    public UpdateResult update(Bson query, Bson updates, boolean upsert, boolean multi) {
        UpdateOptions updateOptions = new UpdateOptions().upsert(upsert);
        if (multi) {
            return dbCollection.updateMany(query, updates, updateOptions);
        } else {
            return dbCollection.updateOne(query, updates, updateOptions);
        }
    }

    public BulkWriteResult replace(List<? extends Bson> queries, List<? extends Bson> updates, boolean upsert) {
        if (queries.size() != updates.size()) {
            throw wrongQueryUpdateSize(queries, updates);
        }

        Iterator<? extends Bson> queryIterator = queries.iterator();
        Iterator<? extends Bson> updateIterator = updates.iterator();

        List<WriteModel<Document>> actions = new ArrayList<>(queries.size());
        UpdateOptions updateOptions = new UpdateOptions().upsert(upsert);

        while (queryIterator.hasNext()) {
            Bson query = queryIterator.next();
            Bson update = updateIterator.next();

            actions.add(new ReplaceOneModel<>(query, (Document) update, updateOptions));
        }

        return dbCollection.bulkWrite(actions, new BulkWriteOptions().ordered(false));
    }

    public UpdateResult replace(Bson query, Bson updates, boolean upsert) {
        UpdateOptions updateOptions = new UpdateOptions().upsert(upsert);
        return dbCollection.replaceOne(query, (Document) updates, updateOptions);
    }

    public BulkWriteResult update(List<? extends Bson> documentList, List<? extends Bson> updatesList,
            boolean upsert, boolean multi) {
        if (documentList.size() != updatesList.size()) {
            throw wrongQueryUpdateSize(documentList, updatesList);
        }

        Iterator<? extends Bson> queryIterator = documentList.iterator();
        Iterator<? extends Bson> updateIterator = updatesList.iterator();

        List<WriteModel<Document>> actions = new ArrayList<>(documentList.size());
        UpdateOptions updateOptions = new UpdateOptions().upsert(upsert);

        while (queryIterator.hasNext()) {
            Bson query = queryIterator.next();
            Bson update = updateIterator.next();

            if (multi) {
                actions.add(new UpdateManyModel<>(query, update, updateOptions));
            } else {
                actions.add(new UpdateOneModel<>(query, update, updateOptions));
            }

            //        BulkWriteOperation bulk = dbCollection.initializeUnorderedBulkOperation();
            //            BulkWriteRequestBuilder builder = bulk.find(query);
            //            if (upsert) {
            //                if (multi) {
            ////                    builder.upsert().update(update);
            //
            //                } else {
            ////                    builder.upsert().updateOne(update);
            //                }
            //            } else {
            //                if (multi) {
            ////                    builder.update(update);
            //                } else {
            ////                    builder.updateOne(update);
            //                }
            //            }

        }
        //        return bulk.execute();
        return dbCollection.bulkWrite(actions, new BulkWriteOptions().ordered(false));
    }

    private IndexOutOfBoundsException wrongQueryUpdateSize(List<? extends Bson> queries,
            List<? extends Bson> updates) {
        return new IndexOutOfBoundsException("QueryList.size=" + queries.size() + " and UpdatesList.size="
                + updates.size() + " must be the same size.");
    }

    public DeleteResult remove(Bson query) {
        return dbCollection.deleteMany(query);
    }

    public DeleteResult remove(Bson query, boolean multi) {
        if (multi) {
            return dbCollection.deleteMany(query);
        } else {
            return dbCollection.deleteOne(query);
        }
    }

    public BulkWriteResult remove(List<? extends Bson> queryList, boolean multi) {
        List<WriteModel<Document>> actions = new ArrayList<>(queryList.size());
        if (multi) {
            for (Bson document : queryList) {
                actions.add(new DeleteManyModel<>(document));
            }
        } else {
            for (Bson document : queryList) {
                actions.add(new DeleteOneModel<>(document));
            }
        }
        return dbCollection.bulkWrite(actions, new BulkWriteOptions().ordered(false));
    }

    public Document findAndUpdate(Bson query, Bson projection, Bson sort, Bson update, QueryOptions options) {
        boolean upsert = false;
        boolean returnNew = false;

        if (options != null) {
            if (projection == null) {
                projection = getProjection(projection, options);
            }
            upsert = options.getBoolean("upsert", false);
            returnNew = options.getBoolean("returnNew", false);
        }

        FindOneAndUpdateOptions findOneAndUpdateOptions = new FindOneAndUpdateOptions().sort(sort)
                .projection(projection).upsert(upsert)
                .returnDocument(returnNew ? ReturnDocument.AFTER : ReturnDocument.BEFORE);
        return dbCollection.findOneAndUpdate(query, update, findOneAndUpdateOptions);
    }

    public Document findAndModify(Bson query, Bson projection, Bson sort, Document update, QueryOptions options) {
        boolean remove = false;
        boolean upsert = false;
        boolean returnNew = false;

        if (options != null) {
            if (projection == null) {
                projection = getProjection(projection, options);
            }
            remove = options.getBoolean("remove", false);
            upsert = options.getBoolean("upsert", false);
            returnNew = options.getBoolean("returnNew", false);
        }

        if (remove) {
            FindOneAndReplaceOptions findOneAndReplaceOptions = new FindOneAndReplaceOptions().sort(sort)
                    .projection(projection).upsert(upsert)
                    .returnDocument(returnNew ? ReturnDocument.AFTER : ReturnDocument.BEFORE);
            return dbCollection.findOneAndReplace(query, update, findOneAndReplaceOptions);
        } else {
            FindOneAndUpdateOptions findOneAndUpdateOptions = new FindOneAndUpdateOptions().sort(sort)
                    .projection(projection).upsert(upsert)
                    .returnDocument(returnNew ? ReturnDocument.AFTER : ReturnDocument.BEFORE);
            return dbCollection.findOneAndUpdate(query, update, findOneAndUpdateOptions);
        }
        //        return dbCollection.findOneAndUpdate(query, projection, sort, remove, update, returnNew, upsert);
    }

    public void createIndex(Bson keys, IndexOptions options) {
        dbCollection.createIndex(keys, options);
    }

    public List<Document> getIndex() {
        return dbCollection.listIndexes().into(new ArrayList<>());
    }

    public void dropIndex(Bson keys) {
        dbCollection.dropIndex(keys);
    }

}