org.fastmongo.odm.bson.repository.BsonMongoTemplate.java Source code

Java tutorial

Introduction

Here is the source code for org.fastmongo.odm.bson.repository.BsonMongoTemplate.java

Source

/*
 * Copyright (c) 2014 Alexander Gulko <kirhog at gmail dot com>.
 *
 * 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.fastmongo.odm.bson.repository;

import com.mongodb.*;
import org.fastmongo.odm.bson.mapping.core.BsonToDomainConverter;
import org.fastmongo.odm.dbobject.mapping.core.DomainToDbObjectConverter;
import org.fastmongo.odm.document.Entity;
import org.fastmongo.odm.repository.MongoSettings;
import org.fastmongo.odm.repository.MongoTemplateOptions;
import org.fastmongo.odm.repository.index.Index;
import org.fastmongo.odm.repository.projection.Projection;
import org.fastmongo.odm.repository.query.Query;
import org.fastmongo.odm.util.MongoNamingService;
import org.fastmongo.odm.util.MongoUtils;

import javax.annotation.PostConstruct;
import java.util.*;

/**
 * The wrapper around {@link com.mongodb.MongoClient} to retrieve data without converting to intermediate DBObjects.
 *
 * @author Alexander Gulko
 */
@SuppressWarnings("UnusedDeclaration")
public class BsonMongoTemplate {
    private MongoClient mongoClient;
    private MongoTemplateOptions templateOptions;
    private MongoSettings settings;
    private BsonToDomainConverter fromBsonConverter;
    private DomainToDbObjectConverter toConverter;
    private MongoNamingService namingService;

    private DB dataBase;
    private final Map<String, DBCollection> collectionsByName = new HashMap<>();

    @PostConstruct
    public void initialize() {
        if (!isMongoDBAvailable()) {
            return;
        }

        dataBase = mongoClient.getDB(templateOptions.getDataBase());

        initializeCollections();
    }

    /**
     * Forces creation of an index on a set of fields, if one does not already exist.
     *
     * @param clazz the collection class.
     * @param keys  pairs with the name of the field or fields to index and order of the index
     */
    public void createIndex(final Class<?> clazz, final Index keys) {
        requestWithDbObject(new Runnable() {
            @Override
            public void run() {
                getCollection(clazz).createIndex(keys.toDbObject());
            }
        });
    }

    public <T> T findById(Class<? extends T> clazz, Object id) {
        String collectionName = getCollectionName(clazz);

        byte[] data = doFindById(collectionName, id);

        return fromBsonConverter.fromBson(clazz, data);
    }

    private byte[] doFindById(String collectionName, Object id) {
        DBCollection collection = getCollection(collectionName);
        DBObject db = collection.findOne(new BasicDBObject(MongoUtils.OBJECT_ID_KEY, id));
        if (db == null) {
            return null;
        }

        return getData(db);
    }

    public <T> List<T> findByIds(Class<? extends T> clazz, Collection<?> ids) {
        return doFindByIds(clazz, ids);
    }

    private <T> List<T> doFindByIds(Class<? extends T> clazz, Collection<?> ids) {
        List<T> result = new ArrayList<>(ids != null ? ids.size() : 0);
        if (ids != null && !ids.isEmpty()) {
            DBCollection collection = getCollection(clazz);
            for (DBObject db : collection.find(new Query().field(MongoUtils.OBJECT_ID_KEY).in(ids).toDbObject())) {
                byte[] data = getData(db);
                T obj = fromBsonConverter.fromBson(clazz, data);
                result.add(obj);
            }
        }

        return result;
    }

    public <T> List<T> find(Class<? extends T> clazz, Query query) {
        return find(clazz, query, null);
    }

    public <T> List<T> find(Class<? extends T> clazz, Query query, Query fields) {
        return doFind(clazz, query, fields);
    }

    private <T> List<T> doFind(Class<? extends T> clazz, Query query, Query fields) {
        List<T> result = new ArrayList<>();
        DBCollection collection = getCollection(clazz);
        for (DBObject db : collection.find(query.toDbObject(), fields != null ? fields.toDbObject() : null)) {
            byte[] data = getData(db);
            T obj = fromBsonConverter.fromBson(clazz, data);
            result.add(obj);
        }

        return result;
    }

    public <T> T findOne(Class<? extends T> clazz, Query query) {
        return findOne(clazz, query, null);
    }

    public <T> T findOne(Class<? extends T> clazz, Query query, Query fields) {
        String collectionName = getCollectionName(clazz);
        byte[] data = doFindOne(collectionName, query, fields);

        return fromBsonConverter.fromBson(clazz, data);
    }

    private byte[] doFindOne(String collectionName, Query query, Query fields) {
        DBCollection collection = getCollection(collectionName);
        DBObject db = collection.findOne(query.toDbObject(), fields != null ? fields.toDbObject() : null);
        if (db == null) {
            return null;
        }

        return getData(db);
    }

    public <T> T findOne(Class<? extends T> clazz, Query query, Query fields, Query orderBy) {
        DBCollection collection = getCollection(clazz);
        DBObject db = collection.findOne(query.toDbObject(), fields != null ? fields.toDbObject() : null,
                orderBy != null ? orderBy.toDbObject() : null);
        if (db == null) {
            return null;
        }

        return read(clazz, db);
    }

    private <T> T read(Class<? extends T> clazz, DBObject db) {
        byte[] data = getData(db);
        return fromBsonConverter.fromBson(clazz, data);
    }

    private <T> T read(Class<? extends T> clazz, byte[] data) {
        return fromBsonConverter.fromBson(clazz, data);
    }

    private byte[] getData(DBObject db) {
        return ((BsonDbObject) db).getDataAndClear();
    }

    /**
     * Gets the count of documents in collection that would match a criteria.
     *
     * @param clazz the collection class.
     * @param query specifies the selection criteria.
     * @return the number of documents that matches selection criteria.
     */
    public long count(Class<?> clazz, Query query) {
        DBCollection collection = getCollection(clazz);
        return collection.count(query != null ? query.toDbObject() : null);
    }

    /**
     * Finds the distinct values for a specified field across a collection and returns the results in a List.
     *
     * @param clazz the collection class.
     * @param key   Specifies the field for which to return the distinct values.
     * @param query specifies the selection query to determine the subset of documents
     *              from which to retrieve the distinct values.
     * @return A {@code List} of the distinct values
     */
    @SuppressWarnings("unchecked")
    public <T> List<T> distinct(Class<?> clazz, String key, Query query) {
        DBCollection collection = getCollection(clazz);
        return collection.distinct(key, query != null ? query.toDbObject() : null);
    }

    /**
     * Saves the given document into its MongoDB collection.
     *
     * @param object the document.
     * @param <T>    the type of the document.
     * @return the saved document.
     */
    public <T extends Entity> T save(T object) {
        return save(object, object.getId());
    }

    /**
     * Saves the given document into its MongoDB collection.
     *
     * @param object        the document.
     * @param mongoObjectId the Mongo _id for the document.
     * @param <T>           the type of the document.
     * @return the saved document.
     */
    public <T> T save(T object, Object mongoObjectId) {
        DBObject db = toConverter.toDbObject(object);
        db.put(MongoUtils.OBJECT_ID_KEY, mongoObjectId);

        DBCollection collection = getCollection(object.getClass());
        collection.save(db);

        return object;
    }

    /**
     * Removes all document from the given collection satisfying the given Query.
     *
     * @param clazz the collection class.
     * @param query specifies the selection criteria (<tt>nullable</tt>).
     */
    public void remove(Class<?> clazz, Query query) {
        DBCollection collection = getCollection(clazz);
        collection.remove(query != null ? query.toDbObject() : null);
    }

    /**
     * Returns cursor to iterate through {@link com.mongodb.BasicDBObject}s.
     *
     * @param clazz      the collection class.
     * @param query      specifies the selection criteria (<tt>nullable</tt>).
     * @param projection specified the fields of document to include/exclude.
     * @return cursor to iterate through parsed {@link com.mongodb.BasicDBObject}s.
     */
    public Iterable<DBObject> findAsDbObjects(final Class<?> clazz, final Query query,
            final Projection projection) {
        DBCollection collection = getCollection(clazz);
        DBCursor cursor = collection.find(query != null ? query.toDbObject() : null,
                projection != null ? projection.toDbObject() : null);
        return new DbObjectCursor(cursor);
    }

    /**
     * Returns a {@link com.mongodb.BasicDBObject}.
     *
     * @param clazz      the collection class.
     * @param query      specifies the selection criteria (<tt>nullable</tt>).
     * @param projection specified the fields of document to include/exclude.
     * @return the {@link com.mongodb.BasicDBObject}.
     */
    public DBObject findAsDbObject(final Class<?> clazz, final Query query, final Projection projection) {
        DBCollection collection = getCollection(clazz);
        DBObject db = collection.findOne(query != null ? query.toDbObject() : null,
                projection != null ? projection.toDbObject() : null);
        return toDbObject(db);
    }

    /**
     * Returns the corresponding {@link com.mongodb.DBCollection} for the given class.
     *
     * @param collectionClass the collection class
     * @return the corresponding {@link com.mongodb.DBCollection} for the given class.
     */
    public DBCollection getCollection(Class<?> collectionClass) {
        return collectionsByName.get(getCollectionName(collectionClass));
    }

    private DBCollection getCollection(String collectionName) {
        return collectionsByName.get(collectionName);
    }

    public String getCollectionName(Class<?> collectionClass) {
        return namingService.getCollectionName(collectionClass);
    }

    public boolean isMongoDBAvailable() {
        return mongoClient != null;
    }

    private void initializeCollections() {
        for (Class<?> collectionClass : templateOptions.getCollectionClasses()) {
            String collectionName = getCollectionName(collectionClass);
            DBCollection collection = dataBase.getCollection(collectionName);

            collectionsByName.put(collectionName, collection);
        }
    }

    private static DBObject toDbObject(DBObject db) {
        if (db instanceof BsonDbObject) {
            db = ((BsonDbObject) db).toDbObjectAndClear();
        }

        return db;
    }

    /**
     * Performs a request to Mongo and switch the template to work with parsed BasicDBObjects.
     *
     * @param task the task to perform.
     */
    public void requestWithDbObject(Runnable task) {
        BsonContextHolder.setUseBson(false);
        try {
            task.run();
        } finally {
            BsonContextHolder.setUseBson(true);
        }
    }

    /**
     * Callback with a result.
     * Differs from {@link java.util.concurrent.Callable} interface in lack of Exceptions in method declaration.
     *
     * @param <T> the type of a result.
     */
    public static interface Task<T> {
        /**
         * Computes a result.
         *
         * @return the result computed by this task.
         */
        T call();
    }

    /**
     * A wrapper for {@link com.mongodb.DBCursor} to retrieve {@link com.mongodb.BasicDBObject}s.
     */
    private static class DbObjectCursor implements Iterable<DBObject> {
        private DBCursor cursor;

        private DbObjectCursor(DBCursor cursor) {
            this.cursor = cursor;
        }

        @Override
        public Iterator<DBObject> iterator() {
            return new DbObjectIterator(cursor.iterator());
        }
    }

    /**
     * A wrapper for {@link com.mongodb.DBCursor} iterator to retrieve {@link com.mongodb.BasicDBObject}s.
     */
    private static class DbObjectIterator implements Iterator<DBObject> {
        private Iterator<DBObject> iterator;

        private DbObjectIterator(Iterator<DBObject> iterator) {
            this.iterator = iterator;
        }

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

        @Override
        public DBObject next() {
            return toDbObject(iterator.next());
        }

        @Override
        public void remove() {
            throw new UnsupportedOperationException();
        }
    }

    public void setFromBsonConverter(BsonToDomainConverter fromBsonConverter) {
        this.fromBsonConverter = fromBsonConverter;
    }

    public void setNamingService(MongoNamingService namingService) {
        this.namingService = namingService;
    }

    public void setMongoClient(MongoClient mongoClient) {
        this.mongoClient = mongoClient;
    }

    public void setTemplateOptions(MongoTemplateOptions templateOptions) {
        this.templateOptions = templateOptions;
    }

    public void setToConverter(DomainToDbObjectConverter toConverter) {
        this.toConverter = toConverter;
    }
}