org.fastmongo.odm.repository.MongoTemplate.java Source code

Java tutorial

Introduction

Here is the source code for org.fastmongo.odm.repository.MongoTemplate.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.repository;

import com.google.common.collect.Lists;
import com.mongodb.*;
import org.fastmongo.odm.repository.projection.Projection;
import org.fastmongo.odm.repository.query.AggregationQuery;
import org.fastmongo.odm.repository.query.Query;
import org.fastmongo.odm.util.MongoNamingService;

import java.util.*;
import java.util.concurrent.TimeUnit;

import static org.fastmongo.odm.util.MongoUtils.OBJECT_ID_KEY;

/**
 * Implementation of {@link MongoOperations} interface.
 *
 * @author Alexander Gulko
 */
public class MongoTemplate implements MongoOperations {
    private MongoClient mongoClient;

    private String dataBase;
    private int maxTimeMS;

    private DB db;
    private MongoTemplateOptions options;
    private final Map<String, DBCollection> collections = new HashMap<>();

    private MongoNamingService namingService;

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

        db = mongoClient.getDB(dataBase);

        for (Class<?> collectionClass : options.getCollectionClasses()) {
            String collectionName = getCollectionName(collectionClass);
            collections.put(collectionName, db.getCollection(collectionName));
        }
    }

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

    public void setOptions(MongoTemplateOptions options) {
        this.options = options;
        this.dataBase = options.getDataBase();
        this.maxTimeMS = options.getMaxTimeMS();
    }

    @Override
    public MongoTemplateOptions getOptions() {
        return options;
    }

    @Override
    public DB getDB() {
        checkIfMongoAvailable();
        return db;
    }

    @Override
    public DBCollection getCollection(String collectionName) {
        checkIfMongoAvailable();
        return collections.get(collectionName);
    }

    @Override
    public DBCollection getCollection(Class<?> collectionClass) {
        return getCollection(getCollectionName(collectionClass));
    }

    @Override
    public Iterable<DBObject> aggregate(Class<?> collectionClass, AggregationQuery query) {
        return aggregate(getCollectionName(collectionClass), query);
    }

    @Override
    public Iterable<DBObject> aggregate(String collectionName, AggregationQuery query) {
        return getResults(getCollection(collectionName).aggregate(query.getPipeline(), aggregationOptions()));
    }

    @Override
    public DBObject findById(Class<?> collectionClass, Object id) {
        return findById(getCollectionName(collectionClass), id);
    }

    @Override
    public DBObject findById(String collectionName, Object id) {
        return doFindById(collectionName, id);
    }

    private DBObject doFindById(String collectionName, Object id) {
        return findOne(collectionName, new Query().field(OBJECT_ID_KEY).eq(id), null);
    }

    @Override
    public List<DBObject> findByIds(Class<?> collectionClass, Collection<?> ids) {
        return findByIds(getCollectionName(collectionClass), ids);
    }

    @Override
    public List<DBObject> findByIds(String collectionName, Collection<?> ids) {
        return doFindByIds(collectionName, ids);
    }

    @Override
    public DBObject findOne(Class<?> collectionClass, Query query, Projection fields) {
        return findOne(collectionClass, query, fields, null);
    }

    @Override
    public DBObject findOne(String collectionName, Query query, Projection fields) {
        return findOne(collectionName, query, fields, null);
    }

    @Override
    public DBObject findOne(Class<?> collectionClass, Query query, Projection fields, Query orderBy) {
        return findOne(getCollectionName(collectionClass), query, fields, orderBy);
    }

    @Override
    public DBObject findOne(String collectionName, Query query, Projection fields, Query orderBy) {
        DBCollection collection = getCollection(collectionName);

        DBObject dbQuery = query.toDbObject();
        DBObject dbFields = fields != null ? fields.toDbObject() : null;
        DBObject dbOrderBy = orderBy != null ? orderBy.toDbObject() : null;

        return collection.findOne(dbQuery, dbFields, dbOrderBy);
    }

    @Override
    public DBCursor find(String collectionName, Query query) {
        DBCollection collection = getCollection(collectionName);
        return configure(collection.find(query.toDbObject()));
    }

    @Override
    public DBCursor find(Class<?> collectionClass, Query query) {
        return find(collectionClass, query, null);
    }

    @Override
    public DBCursor find(Class<?> collectionClass, Query match, Projection fields) {
        return find(getCollectionName(collectionClass), match, fields);
    }

    @Override
    public DBCursor find(String collectionName, Query match, Projection fields) {
        DBCollection collection = getCollection(collectionName);
        return configure(collection.find(match.toDbObject(), fields != null ? fields.toDbObject() : null));
    }

    @Override
    public long count(Class<?> collectionClass, Query query) {
        DBCollection collection = getCollection(collectionClass);
        return collection.count(query.toDbObject());
    }

    @Override
    public WriteResult update(String collectionName, Query query, DBObject o) {
        DBCollection collection = getCollection(collectionName);
        return collection.update(query.toDbObject(), o);
    }

    @Override
    @SuppressWarnings("unchecked")
    public <T> List<T> distinct(String collectionName, String key, Query query) {
        DBCollection collection = getCollection(collectionName);
        return collection.distinct(key, query.toDbObject());
    }

    @SuppressWarnings("unchecked")
    @Override
    public <T> List<T> distinct(Class<?> collectionClass, String key, Query query) {
        DBCollection collection = getCollection(getCollectionName(collectionClass));
        return collection.distinct(key, query.toDbObject());
    }

    @Override
    public void save(Class<?> collectionClass, DBObject dbObject) {
        getCollection(getCollectionName(collectionClass)).save(dbObject);
    }

    @Override
    public void save(String collectionName, DBObject dbObject) {
        getCollection(collectionName).save(dbObject);
    }

    @Override
    public BulkWriteResult save(Class<?> collectionClass, Collection<DBObject> dbObjects) {
        BulkWriteOperation bulk = getCollection(collectionClass).initializeUnorderedBulkOperation();

        for (DBObject dbObject : dbObjects) {
            DBObject idQuery = new BasicDBObject();
            idQuery.put(OBJECT_ID_KEY, dbObject.get(OBJECT_ID_KEY));
            bulk.find(idQuery).upsert().replaceOne(dbObject);
        }

        return bulk.execute();
    }

    @Override
    public void createIndex(String collectionName, Query keys) {
        getCollection(collectionName).createIndex(keys.toDbObject());
    }

    @Override
    public void remove(String collectionName, DBObject dbObject) {
        getCollection(collectionName).findAndRemove(dbObject);
    }

    @Override
    public void remove(Class<?> collectionClass, DBObject dbObject) {
        getCollection(getCollectionName(collectionClass)).remove(dbObject);
    }

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

    private void checkIfMongoAvailable() {
        if (!isMongoDBAvailable()) {
            throw new IllegalStateException("MongoClient's method is called but MongoDB is disabled. "
                    + "Please ensure that the 'mongo.enabled' flag is set to true, "
                    + "or 'nosql.enabled' is set to 'mongo'.");
        }
    }

    List<DBObject> doFindByIds(String collectionName, Collection<?> ids) {
        List<DBObject> result = Collections.emptyList();
        if (!ids.isEmpty()) {
            result = new ArrayList<>(ids.size());
            DBCursor cursor = find(collectionName, new Query().field(OBJECT_ID_KEY).in(ids), null);
            for (DBObject dbObject : cursor) {
                result.add(dbObject);
            }
        }
        return result;
    }

    private DBCursor configure(DBCursor dbCursor) {
        return dbCursor.maxTime(maxTimeMS, TimeUnit.MILLISECONDS);
    }

    private AggregationOptions aggregationOptions() {
        return AggregationOptions.builder().maxTime(maxTimeMS, TimeUnit.MILLISECONDS).build();
    }

    private Iterable<DBObject> getResults(Cursor cursor) {
        return Lists.newArrayList(cursor);
    }

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

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