D-Net

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

import java.util.*;

import java.util.concurrent.*;

import java.util.regex.Pattern;

import com.google.common.base.Function;

import com.google.common.collect.Iterables;

import com.google.common.collect.Iterators;

import com.google.common.collect.Lists;

import com.google.common.collect.Sets;

import com.mongodb.BasicDBObject;

import com.mongodb.DBObject;

import com.mongodb.QueryBuilder;

import com.mongodb.client.FindIterable;

import com.mongodb.client.ListIndexesIterable;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.MongoDatabase;

import com.mongodb.client.model.IndexOptions;

import eu.dnetlib.data.mdstore.DocumentNotFoundException;

import eu.dnetlib.data.mdstore.MDStoreServiceException;

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.enabling.resultset.ResultSetListener;

import eu.dnetlib.enabling.resultset.client.ResultSetClient;

import org.apache.commons.lang3.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.conversions.Bson;

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

public class MongoMDStore implements MDStore {

	private static final int BULK_SIZE = 500;

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

	private final boolean discardRecords;

	private String id;

	private MongoDatabase mongoDatabase;

	private MongoCollection<DBObject> collection;

	private MongoCollection<DBObject> discardedCollection;

	private List<MDFormatDescription> mdformats;

	private RecordParser recordParser;

	private static List<String> indices = Lists.newArrayList("id", "timestamp", "originalId");

	private final IndexOptions options = new IndexOptions().background(true);

	public MongoMDStore(final String id,

			final MongoCollection<DBObject> collection,

			final RecordParser recordParser,

			final boolean discardRecords,

			final MongoDatabase mongoDatabase) {

		this.id = id;

		this.mongoDatabase = mongoDatabase;

		this.collection = collection;

		this.discardedCollection = this.mongoDatabase.getCollection("discarded-" + StringUtils.substringBefore(id, "_"), DBObject.class);

		this.recordParser = recordParser;

		this.discardRecords = discardRecords;

	}

	@Override

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

		this.mdformats = mdformats;

		return feed(records, incremental);

	}

	@Override

	public int 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<>(100);

		final Object sentinel = new Object();

		int countStored = 0;

		final Callable<Integer> writer = () -> {

			final MongoBulkWritesManager bulkWritesManager =

					new MongoBulkWritesManager(collection, discardedCollection, mdformats, BULK_SIZE, recordParser, discardRecords);

			int count = 0;

			while (true) {

				try {

					final Object record = queue.take();

					if(record == ResultSetClient.SENTINEL_PROC_CANCELED){

						//it means the wf was cancelled, so we must give this information back to the caller to avoid commit

						log.warn(String.format("Process cancelled after %d records, returning -1", count));

						return -1;

					}

					if (record == sentinel) {

						bulkWritesManager.flushBulks();

						break;

					}

					count++;

					bulkWritesManager.insert((String) record);

				} catch (final InterruptedException e) {

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

					throw new IllegalStateException(e);

				}

			}

			log.debug(String.format("extracted %s records from feeder queue", count));

			return count;

		};

		final ExecutorService executorService = Executors.newSingleThreadExecutor();

		Future<Integer> storedCountInt = executorService.submit(writer);

		try {

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

			if (records != null) {

				for (final String record : records) {

					queue.put(record);

				}

			}

			queue.put(sentinel);

			countStored = storedCountInt.get().intValue();

		} catch (final InterruptedException e) {

			log.error("Error on feeding mdstore with id:" + id, e);

			throw new IllegalStateException(e);

		} catch (ExecutionException e) {

			log.error("Error on feeding mdstore with id:" + id, e);

			throw new IllegalStateException(e);

		}

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

		return countStored;

	}

	public void ensureIndices() {

		for (final String key : indices) {

			collection.createIndex(new BasicDBObject(key, 1), options);

		}

		if (mdformats != null) {

			for (final MDFormatDescription description : mdformats) {

				collection.createIndex(new BasicDBObject(description.getName(), 1), options);

			}

		}

	}

	public boolean isIndexed() {

		final ListIndexesIterable<DBObject> found = collection.listIndexes(DBObject.class);

		return Sets.newHashSet(Iterables.transform(found, dbo -> {

			final Set<String> keyset = ((DBObject) dbo.get("key")).toMap().keySet();

			return Iterables.getFirst(keyset, "");

		})).containsAll(indices);

	}

	/**

	 * 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);

		BasicDBObject query = (BasicDBObject) QueryBuilder.start("body").regex(regex).get();

		final FindIterable<DBObject> matches = collection.find(query, DBObject.class);

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

		if (log.isDebugEnabled())

			log.debug("FOUND: " + Lists.newArrayList(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.findOneAndReplace(new BasicDBObject("_id", o.get("_id")), o);

		}

	}

	@Override

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

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

	}

	@Override

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

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

	}

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

			throws MDStoreServiceException {

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

		return new MongoResultSetListener(collection, parseLong(from), parseLong(until), filter, serializer);	}

	private Long parseLong(final String s) throws MDStoreServiceException {

		if (StringUtils.isBlank(s)) {

			return null;

		}

		try {

			return Long.valueOf(s);

		} catch (NumberFormatException e) {

			throw new MDStoreServiceException("Invalid date, expected java.lang.Long, or null", e);

		}

	}

	@Override

	public Iterable<String> iterate() {

		return () -> Iterators.transform(collection.find().iterator(), arg -> (String) arg.get("body"));

	}

	@Override

	public void deleteRecord(final String recordId) {

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

	}

	@Override

	public String getRecord(final String recordId) throws DocumentNotFoundException {

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

		if (obj == null || !obj.containsField("body")) 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().apply(obj);

	}

	@Override

	public List<String> deliver(final String mdId, final int pageSize, final int offset, final Map<String, String> queryParam) {

		final QueryBuilder query = QueryBuilder.start();

		for (String key : queryParam.keySet()) {

			query.and(key).regex(Pattern.compile(queryParam.get(key), Pattern.LITERAL));

		}

		FindIterable<DBObject> dbObjects = offset > 0

				? collection.find((Bson) query.get()).limit(pageSize).skip(offset)

				: collection.find((Bson) query.get()).limit(pageSize);

		queryParam.put("count", "" + collection.count((Bson) query.get()));

		final List<String> result = new ArrayList<>();

		for (final DBObject item : dbObjects) {

			result.add(item.get("body").toString());

		}

		return result;

	}

	@Override

	public void truncate() {

		collection.drop();

		discardedCollection.drop();

	}

	public DBObject getMDStoreMetadata() {

		return mongoDatabase.getCollection("metadata", DBObject.class).find(new BasicDBObject("mdId", getId())).first();

	}

	@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 MongoCollection<DBObject> getCollection() {

		return collection;

	}

	public void setCollection(final MongoCollection<DBObject> 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.count();

	}

	public MongoCollection<DBObject> getDiscardedCollection() {

		return discardedCollection;

	}

	public void setDiscardedCollection(final MongoCollection<DBObject> discardedCollection) {

		this.discardedCollection = discardedCollection;

	}

	private class SerializeMongoRecord implements Function<DBObject, String> {

		@Override

		public String apply(final DBObject arg) {

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

		}

	}

	private class SerializeMongoRecordId implements Function<DBObject, String> {

		@Override

		public String apply(final DBObject arg) {

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

		}

	}

}