com.thematchbox.db.MongoDataSource Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of TheMatchBox-Common Show documentation
This project contains some common utilities used in different projects from theMatchBox.
There is a newer version: 2.0.2
package com.thematchbox.db;

/* Copyright 2015 theMatchBox

   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.
*/


import com.mongodb.*;
import org.bson.types.ObjectId;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;

public class MongoDataSource {

    public static final Logger logger = LoggerFactory.getLogger(MongoDataSource.class);

    private ObjectFactory objectFactory;
    private ObjectCache objectCache;

    private String collectionName;

    public MongoDataSource(ObjectFactory objectFactory, ObjectCache objectCache) {
        this.objectFactory = objectFactory;
        this.collectionName = objectFactory.name();
        this.objectCache = objectCache;
    }

    public String getCollectionName() {
        return collectionName;
    }

    public ObjectFactory getObjectFactory() {
        return objectFactory;
    }

    public D read(ObjectId objectId, MongoDBConnection dbConnection) throws MongoDBException {
        if (objectCache.containsKey(objectId)) {
            return objectCache.get(objectId);
        }

        DBCollection collection = dbConnection.db.getCollection(collectionName);
        DBObject one = collection.findOne(new BasicDBObject(DataObject.ID_FIELD, objectId));
        D d = one != null ? objectFactory.unmarshall(one) : null;
        objectCache.put(objectId, d);
        return d;
    }

    public void write(D d, MongoDBConnection dbConnection) throws MongoDBException {
        DBObject dbObject = objectFactory.marshall(d, dbConnection);
        objectCache.put(d.getObjectId(), d);
        DBCollection collection = dbConnection.db.getCollection(collectionName);
        collection.update(new BasicDBObject(DataObject.ID_FIELD, d.getObjectId()), dbObject, true, false);
        d.setStored(true);
    }

    public void delete(ObjectId objectId, MongoDBConnection dbConnection) throws MongoDBException {
        DBCollection collection = dbConnection.db.getCollection(collectionName);
        collection.remove(new BasicDBObject(DataObject.ID_FIELD, objectId));
    }

    public LazyReference findOne(DBObject query, MongoDBConnection dbConnection) {

        DBCollection collection = dbConnection.db.getCollection(collectionName);
        DBObject dbObject = collection.findOne(query, new BasicDBObject(DataObject.ID_FIELD, 1));
        if (dbObject != null) {
            ObjectId objectId = objectFactory.get(dbObject, DataObject.ID_FIELD);

            return objectCache.containsKey(objectId) ? new LazyReference<>(objectCache.get(objectId)) : new LazyReference(objectId);
        }
        return null;
    }

    public MongoDBIterator findObjects(DBObject query, MongoDBConnection dbConnection) {
        DBCollection collection = dbConnection.db.getCollection(collectionName);
        final DBCursor cursor = collection.find(query);
        return new MongoDBIterator<>(cursor, objectFactory, objectCache);

    }

    public PageResult findReferences(DBObject query, int from, int batchSize, MongoDBConnection dbConnection) {
        return findReferences(query, from, batchSize, null, dbConnection);
    }

    public PageResult findReferences(DBObject query, int from, int batchSize, DBObject orderBy, MongoDBConnection dbConnection) {
        DBCollection collection = dbConnection.db.getCollection(collectionName);
        final DBCursor cursor = collection.find(query);
        if (orderBy != null && !orderBy.keySet().isEmpty()) {
            cursor.sort(orderBy);
        }
        cursor.skip(from);
        cursor.limit(batchSize);
        int count = cursor.count();
        List> references = new ArrayList<>(batchSize);
        while (cursor.hasNext()) {
            DBObject dbObject = cursor.next();
            ObjectId objectId = objectFactory.get(dbObject, DataObject.ID_FIELD);
            references.add(new LazyReference(objectId));
        }
        return new PageResult<>(references, from, count);

    }

    public MongoDBIterator getAllObjects(MongoDBConnection dbConnection) {
        DBCollection collection = dbConnection.db.getCollection(collectionName);
        final DBCursor cursor = collection.find();
        return new MongoDBIterator<>(cursor, objectFactory, objectCache);
    }


    public List> getAllReferences(MongoDBConnection dbConnection) {
        return findReferences(new BasicDBObject(), dbConnection);
    }

    public List> findReferences(BasicDBObject query, MongoDBConnection dbConnection) {
        List> references = new ArrayList<>();
        DBCollection collection = dbConnection.db.getCollection(collectionName);
        final DBCursor cursor = collection.find(query, new BasicDBObject(DataObject.ID_FIELD, 1));
        while (cursor.hasNext()) {
            DBObject dbObject = cursor.next();
            ObjectId objectId = objectFactory.get(dbObject, DataObject.ID_FIELD);
            references.add(new LazyReference(objectId));
        }
        return references;
    }

    public MongoDBIterator getObjectsForIds(MongoDBConnection dbConnection, List ids) {
        BasicDBList dbList = new BasicDBList();
        for (ObjectId id : ids) {
            dbList.add(id);
        }
        BasicDBObject query = new BasicDBObject(DataObject.ID_FIELD, new BasicDBObject("$in", dbList));

        DBCollection collection = dbConnection.db.getCollection(collectionName);
        final DBCursor cursor = collection.find(query);
        return new MongoDBIterator<>(cursor, objectFactory, objectCache);
    }



    public void update(D object, DBObject query, DBObject update, MongoDBConnection dbConnection) {
        DBCollection collection = dbConnection.db.getCollection(getCollectionName());
        objectCache.put(object.getObjectId(), object);
        collection.update(query, update, false, false);
    }

    public void ensureIndex(MongoDBConnection dbConnection, String path) {
        DBCollection collection = dbConnection.db.getCollection(collectionName);
        collection.ensureIndex(path);
    }

    public String buildPath(String... parts) {
        StringBuilder builder = new StringBuilder();
        for (String part : parts) {
            if (builder.length() > 0) {
                builder.append('.');
            }
            builder.append(part);
        }

        return builder.toString();
    }

    public void drop(MongoDBConnection dbConnection) {
        DBCollection collection = dbConnection.db.getCollection(collectionName);
        collection.drop();
    }

}