D-Net

package eu.dnetlib.data.mdstore.modular.mongodb;

import java.util.Iterator;

import java.util.List;

import java.util.concurrent.ArrayBlockingQueue;

import java.util.concurrent.BlockingQueue;

import java.util.regex.Pattern;

import org.apache.commons.lang.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.BasicBSONObject;

import org.springframework.beans.factory.annotation.Required;

import com.google.common.collect.Iterators;

import com.google.common.collect.Lists;

import com.mongodb.BasicDBObject;

import com.mongodb.DBCollection;

import com.mongodb.DBCursor;

import com.mongodb.DBObject;

import com.mongodb.QueryBuilder;

import eu.dnetlib.data.mdstore.DocumentNotFoundException;

import eu.dnetlib.data.mdstore.modular.MDFormatDescription;

import eu.dnetlib.data.mdstore.modular.RecordParser;

import eu.dnetlib.data.mdstore.modular.connector.MDStore;

import eu.dnetlib.data.mdstore.modular.mongodb.utils.MDStoreUtils;

import eu.dnetlib.enabling.resultset.ResultSetListener;

import eu.dnetlib.miscutils.collections.MappedCollection;

import eu.dnetlib.miscutils.functional.UnaryFunction;

public class MongoMDStore implements MDStore {

	private static final int BULK_SIZE = 500;

	private static final class SerializeMongoRecord implements UnaryFunction<String, DBObject> {

		@Override

		public String evaluate(final DBObject arg) {

			return (String) arg.get("body");

		}

	}

	private static final class SerializeMongoRecordId implements UnaryFunction<String, DBObject> {

		@Override

		public String evaluate(final DBObject arg) {

			return (String) arg.get("id");

		}

	}

	private static final Log log = LogFactory.getLog(MongoMDStore.class);

	private String id;

	private DBCollection collection;

	private DBCollection discardedCollection;

	private RecordParser recordParser;

	private final boolean discardRecords;

	private static List<String> requiredIndicies = Lists.newArrayList("{ \"id\" : 1}", "{ \"timestamp\" : 1}", "{ \"originalId\" : 1}");

	public MongoMDStore(final String id, final DBCollection collection, final RecordParser recordParser, final boolean discardRecords) {

		this.id = id;

		this.collection = collection;

		this.discardedCollection = collection.getDB().getCollection("discarded-" + StringUtils.substringBefore(id, "_"));

		this.recordParser = recordParser;

		this.discardRecords = discardRecords;

	}

	@Override

	public void feed(final Iterable<String> records, final boolean incremental) {

		// TODO: remove incremental from MDStore API. It is used in MDStoreModular. Useless here.

		ensureIndices();

		final BlockingQueue<Object> queue = new ArrayBlockingQueue<Object>(100);

		final Object sentinel = new Object();

		final Thread background = new Thread(new Runnable() {

			@Override

			public void run() {

				MongoBulkWritesManager bulkWritesManager = new MongoBulkWritesManager(collection, discardedCollection, BULK_SIZE, recordParser, discardRecords);

				while (true) {

					try {

						final Object record = queue.take();

						if (record == sentinel) {

							bulkWritesManager.flushBulks();

							break;

						}

						bulkWritesManager.insert((String) record);

					} catch (final InterruptedException e) {

						log.fatal("got exception in background thread", e);

						throw new IllegalStateException(e);

					}

				}

			}

		});

		background.start();

		try {

			log.info("feeding mdstore " + id);

			if (records != null) {

				for (final String record : records) {

					queue.put(record);

				}

			}

			queue.put(sentinel);

			log.info("finished feeding mdstore " + id);

			background.join();

		} catch (final InterruptedException e) {

			throw new IllegalStateException(e);

		}

		// double check

		ensureIndices();

		collection.ensureIndex(new BasicDBObject("id", 1));

	}

	public void ensureIndices() {

		collection.resetIndexCache();

		for (final String key : Lists.newArrayList("id", "timestamp", "originalId")) {

			collection.ensureIndex(new BasicDBObject(key, 1));

		}

	}

	public boolean isIndexed() {

		return Lists.newArrayList(new MappedCollection<String, DBObject>(collection.getIndexInfo(), new UnaryFunction<String, DBObject>() {

			@Override

			public String evaluate(final DBObject dbo) {

				return new BasicBSONObject(dbo.toMap()).getString("key");

			}

		})).containsAll(requiredIndicies);

	}

	/**

	 * Method searches for the given string grep into this collection and replaces it with the given replacement.

	 *

	 * @param grep

	 *            the string to search

	 * @param replace

	 *            the replacement

	 */

	public void replace(final String grep, final String replace) {

		final Pattern regex = Pattern.compile(grep, Pattern.MULTILINE);

		final DBCursor matches = collection.find(QueryBuilder.start("body").regex(regex).get());

		log.debug("FOUND: " + matches.size());

		for (final DBObject match : matches) {

			final DBObject o = new BasicDBObject(match.toMap());

			o.put("body", regex.matcher((String) match.get("body")).replaceAll(replace));

			collection.update(match, o);

		}

	}

	@Override

	public ResultSetListener deliver(final String from, final String until, final String recordFilter) {

		return deliver(from, until, recordFilter, new SerializeMongoRecord());

	}

	@Override

	public ResultSetListener deliverIds(final String from, final String until, final String recordFilter) {

		return deliver(from, until, recordFilter, new SerializeMongoRecordId());

	}

	public ResultSetListener deliver(final String from, final String until, final String recordFilter, final UnaryFunction<String, DBObject> serializer) {

		ensureIndices();

		final Pattern filter = recordFilter != null && recordFilter.length() > 0 ? Pattern.compile(recordFilter, Pattern.MULTILINE) : null;

		return new MongoResultSetListener(collection, filter, serializer);

	}

	@Override

	public Iterable<String> iterate() {

		return new Iterable<String>() {

			@Override

			public Iterator<String> iterator() {

				return Iterators.transform(collection.find(), MDStoreUtils.body());

			}

		};

	}

	@Override

	public void deleteRecord(final String recordId) {

		collection.remove(new BasicDBObject("id", recordId));

	}

	@Override

	public String getRecord(final String recordId) throws DocumentNotFoundException {

		final DBObject obj = collection.findOne(new BasicDBObject("id", recordId));

		if (obj == null || obj.containsField("body") == false) { throw new DocumentNotFoundException(String.format(

				"The document with id '%s' does not exist in mdstore: '%s'", recordId, id)); }

		final String body = (String) obj.get("body");

		if (body.trim().length() == 0) { throw new DocumentNotFoundException(String.format("The document with id '%s' does not exist in mdstore: '%s'",

				recordId, id)); }

		return new SerializeMongoRecord().evaluate(obj);

	}

	@Override

	public void truncate() {

		collection.drop();

		discardedCollection.drop();

	}

	public DBObject getMDStoreMetadata() {

		return collection.getDB().getCollection("metadata").findOne(new BasicDBObject("mdId", getId()));

	}

	@Override

	public String getFormat() {

		return (String) getMDStoreMetadata().get("format");

	}

	@Override

	public String getInterpretation() {

		return (String) getMDStoreMetadata().get("interpretation");

	}

	@Override

	public String getLayout() {

		return (String) getMDStoreMetadata().get("layout");

	}

	@Override

	public String getId() {

		return id;

	}

	public void setId(final String id) {

		this.id = id;

	}

	public DBCollection getCollection() {

		return collection;

	}

	public void setCollection(final DBCollection collection) {

		this.collection = collection;

	}

	public RecordParser getRecordParser() {

		return recordParser;

	}

	@Required

	public void setRecordParser(final RecordParser recordParser) {

		this.recordParser = recordParser;

	}

	@Override

	public int getSize() {

		return (int) collection.getCount();

	}

	public DBCollection getDiscardedCollection() {

		return discardedCollection;

	}

	public void setDiscardedCollection(final DBCollection discardedCollection) {

		this.discardedCollection = discardedCollection;

	}

	@Override

	public void feed(final Iterable<String> records, final boolean incremental, final List<MDFormatDescription> mdformats) {

		// TODO Auto-generated method stub

	}

}