com.setronica.ucs.storage.MongoMapStorage.java Source code

Introduction

Here is the source code for com.setronica.ucs.storage.MongoMapStorage.java

Source

/*
 * Copyright (C) 2014 SETRONICA - setronica.com
 * This source code is available under the terms of the GNU Lesser General Public License
 * as published by The Open Source Initiative (OSI), either version 3 of the License,
 * or (at your option) any later version.
 *
 * UCS is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
 * without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 *
 * See the GNU Lesser General Public License for more details.
 * You should have received a copy of the GNU Lesser General Public License.
 */

package com.setronica.ucs.storage;

import com.mongodb.*;
import org.bson.BSONObject;

import java.io.*;
import java.util.HashMap;

public class MongoMapStorage<K, E> extends MemMapStorage<K, E> implements PersistentStorage {
    private String mongoDbName;
    private String mongoCollectionName;
    protected DBCollection collection;
    protected MongoClient mongoClient;

    public MongoMapStorage(MongoClient mongoClient, String mongoDbName, String mongoCollectionName) {
        this.mongoClient = mongoClient;
        this.mongoDbName = mongoDbName;
        this.mongoCollectionName = mongoCollectionName;
    }

    protected void persist(K key, E value) {
        try {
            BasicDBObject object = new BasicDBObject("_id", writeObject(key));
            object.put("value", writeObject(value));
            collection.save(object);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    protected void delete(K key, E value) {
        try {
            collection.remove(new BasicDBObject("_id", writeObject(key)));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public void open() throws IOException {
        DB db = mongoClient.getDB(mongoDbName);
        collection = db.getCollection(mongoCollectionName);
        if (collection == null) {
            collection = db.createCollection(mongoCollectionName, null);
        }
        collection.setWriteConcern(WriteConcern.ACKNOWLEDGED);

        DBCursor cursor = collection.find();
        data = new HashMap<K, E>(cursor.size());
        for (DBObject entry : cursor) {
            data.put((K) readObject(entry.get("_id")), (E) readObject(entry.get("value")));
        }
    }

    @Override
    public void close() throws IOException {
    }

    private Object writeObject(Object element) throws IOException {
        if (element instanceof BSONObject || element instanceof Number || element instanceof String) {
            return element;
        }
        ByteArrayOutputStream os = new ByteArrayOutputStream();
        ObjectOutput out = new ObjectOutputStream(os);
        out.writeObject(element);
        byte[] val = os.toByteArray();
        out.close();
        os.close();
        return val;
    }

    private Object readObject(Object data) throws IOException {
        if (data instanceof BSONObject || data instanceof Number || data instanceof String) {
            return data;
        } else if (!(data instanceof byte[])) {
            throw new RuntimeException("unexpected object: " + data.getClass().getName());
        }
        ByteArrayInputStream bis = new ByteArrayInputStream((byte[]) data);
        ObjectInput ois = new ObjectInputStream(bis);
        try {
            return ois.readObject();
        } catch (ClassNotFoundException e) {
            throw new IOException(e);
        } finally {
            ois.close();
            bis.close();
        }
    }
}