D-Net

package eu.dnetlib.data.mdstore.mongo;

import java.util.Arrays;

import java.util.List;

import java.util.Map;

import java.util.Set;

import java.util.stream.Collectors;

import java.util.stream.Stream;

import javax.annotation.PostConstruct;

import com.google.common.collect.Lists;

import com.google.common.collect.Sets;

import com.mongodb.BasicDBObject;

import com.mongodb.WriteConcern;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.MongoDatabase;

import com.mongodb.client.model.IndexOptions;

import com.mongodb.client.model.UpdateOptions;

import eu.dnetlib.data.mdstore.MDStore;

import eu.dnetlib.data.mdstore.MDStoreService;

import eu.dnetlib.data.mdstore.model.MDStoreRecord;

import eu.dnetlib.data.mdstore.model.Metadata;

import eu.dnetlib.data.mdstore.model.Transaction;

import eu.dnetlib.data.mdstore.repository.MetadataRepository;

import eu.dnetlib.data.mdstore.utils.MDStoreException;

import eu.dnetlib.data.mdstore.utils.MDStoreProfileManager;

import eu.dnetlib.data.mdstore.utils.MDStoreRuntimeException;

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

import eu.dnetlib.miscutils.streams.DnetStreamSupport;

import org.apache.commons.lang3.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

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

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

import org.springframework.stereotype.Component;

import static com.mongodb.client.model.Filters.eq;

import static eu.dnetlib.data.mdstore.utils.MDStoreConstants.DISCARDED_PREFIX;

import static eu.dnetlib.data.mdstore.utils.MDStoreConstants.ID;

/**

 * Created by claudio on 24/03/2017.

 */

@Component

public class MongoMDStoreService implements MDStoreService {

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

	/**

	 * The max number of concurrent transactions per mdstore.

	 */

	private int maxTransactions = 5;

	/**

	 * The expired days.

	 */

	private int expiredDays = 3;

	@Autowired

	private MDStoreProfileManager profileManager;

	@Autowired

	private MetadataRepository metadata;

	@Autowired

	private MongoDatabase mongoDatabase;

	@Value("#{'${msro.worker.mdstore.skip.prefixes}'.split(',')}")

	private List<String> skipPrefixes = Lists.newArrayList("resolved", "metadata");

	/**

	 * Bootstrap manager.

	 */

	@PostConstruct

	private void init() {

		log.info("Bootstrap mdstore ...");

		//drop all the stale expiring and the transactions mdstores

		metadata.findAll().forEach(d -> {

			metadata.bootstrap(d.getMdId());

			log.debug(String.format("Added %s to Metadata Manager data structure", d.getMdId()));

		});

		final Set<String> metadataMdIds = Sets.newHashSet(metadata.listMdIds());

		//delete from metadataManager all the mdId(s) that are not in metadata

		Sets.difference(

				Sets.newHashSet(metadata.listMdIds()),

				metadataMdIds).forEach(mdId -> metadata.delete(mdId));

		//delete all the collections that are not referred by metadata.currentId(s)

		Sets.difference(

				Sets.newHashSet(metadata.listCurrentIds()),

				metadataMdIds).forEach(name -> {

						log.info(String.format("dropping collection %s", name));

						mongoDatabase.getCollection(name).drop();

					});

		log.info("Bootstrap mdstore complete!");

	}

	@Override

	public String create(final String format, final String layout, final String interpretation) throws MDStoreException {

		log.debug("Creating new mdstore");

		final String mdId = getProfileManager().registerProfile(format, layout, interpretation);

		metadata.save(Metadata.create()

				.setMdId(mdId)

				.setCurrentId(mdId)

				.setFormat(format)

				.setLayout(layout)

				.setInterpretation(interpretation));

		mongoDatabase.createCollection(mdId);

		((MongoMDStore) getMDStore(mdId)).ensureIndices();

		log.info("Created new mdstore: " + mdId);

		return mdId;

	}

	@Override

	public void delete(final String mdId) throws MDStoreException {

		if (!metadata.exists(mdId)) {

			throw new MDStoreException(String.format("cannot find %s in transaction manager table", mdId));

		}

		metadata.delete(mdId);

		getMDStore(mdId).truncate();

		getProfileManager().deleteProfile(mdId);

		log.info("deleted mdId: " + mdId);

	}

	@Override

	public Stream<MDStoreRecord> stream(final String mdId) throws MDStoreException {

		return stream(mdId, null, null);

	}

	@Override

	public Stream<MDStoreRecord> stream(final String mdId, final String from, final String until) throws MDStoreException {

		return stream(mdId, from, until, null);

	}

	@Override

	public Stream<MDStoreRecord> stream(final String mdId, final String from, final String until, final String recordFilter) throws MDStoreException {

		return getMDStore(mdId).deliver(from, until, recordFilter);

	}

	@Override

	public Stream<MDStoreRecord> bulkStream(final String format, final String layout, final String interpretation) throws MDStoreException {

		return list(format, layout, interpretation).stream()

				.map(mdId -> {

					log.debug("bulk deliver of mdId: " + mdId);

					try {

						return getMDStore(mdId).deliver();

					} catch (MDStoreException e) {

						throw new MDStoreRuntimeException(e);

					}

				}).reduce(Stream::concat).orElseGet(Stream::empty);

	}

	@Override

	public MDStoreRecord findOne(final String mdId, final String recordId) throws MDStoreException {

		return getMDStore(mdId).getRecord(recordId);

	}

	@Override

	public List<String> list() throws MDStoreException {

		return DnetStreamSupport.stream(metadata.findAll().iterator())

			.map(m -> m.getMdId())

			.collect(Collectors.toList());

	}

	@Override

	public List<String> list(final String format, final String layout, final String interpretation) throws MDStoreException {

		return metadata.findByFormatAndLayoutAndInterpretation(format, layout, interpretation).stream()

				.map(m -> m.getMdId())

				.collect(Collectors.toList());

	}

	@Override

	public long size(final String mdId) throws MDStoreException {

		return getMDStore(mdId).getSize();

	}

	@Override

	public long size(final String format, final String layout, final String interpretation) {

		return metadata.sumOfSize(format, layout, interpretation);

	}

	@Override

	public void store(final String mdId, final Stream<String> records, final FeedMode feedMode) throws MDStoreException {

		final boolean refresh = FeedMode.REFRESH.equals(feedMode);

		final String txId = startTransaction(mdId, refresh);

		final MDStore mdstore = getMDStore(mdId, txId);

		try {

			if (refresh) {

				mdstore.truncate();

			}

			mdstore.feed(records.map(new RecordParser()));

			final Metadata m = metadata.findByMdId(mdId);

			final Map<String, Transaction> txs = m.getTransactionMap();

			if (!txs.containsKey(txId)) {

				throw new MDStoreException("Error, unable to find transaction with Id " + txId);

			}

			if (txs.get(txId).getRefresh()) {

				if (!m.isCurrentIdExpired(expiredDays)) {

					doDrop(m.getCurrentId());

				}

				m.setCurrentId(txId);

				log.debug(String.format("commit refresh: replaced collection %s with %s", mdId, txId));

			} else {

				updateIncremental(txId, m.getCurrentId());

				doDrop(txId);

				log.debug(String.format("commit incremental: "));

			}

			final long size = getRecordCollection(m.getCurrentId()).count();

			metadata.commit(m.setSize(size), txId);

			getProfileManager().updateSize(mdId, size);

			log.info("Finished feeding mdstore " + mdId + " - new size: " + size);

		} catch (Throwable e) {

			if (txId != null) {

				log.error("Dropping transaction : " + txId);

				metadata.dropTransaction(mdId, txId);

			}

			throw new MDStoreException("Error feeding mdstore: " + mdId, e);

		}

	}

	/// helpers

	private String startTransaction(final String mdId, final boolean refresh) throws MDStoreException {

		log.info("Start transaction for metadata store " + mdId);

		final Metadata m = metadata.findByMdId(mdId);

		if (m.getTransactions().size() > getMaxTransactions()) {

			throw new MDStoreException(

					String.format("Cannot create more than %s transactions, found: %s, mdId: %s", getMaxTransactions(), m.getTransactions().size(), mdId));

		}

		final String txId = mdId + "::" + System.currentTimeMillis();

		metadata.addTransaction(mdId,

				Transaction.create()

					.setId(txId)

					.setRefresh(refresh));

		return txId;

	}

	/**

	 * Update incremental.

	 *

	 * @param transactionId the transaction id

	 * @param currentId     the current id

	 */

	private void updateIncremental(final String transactionId, final String currentId) {

		final MongoCollection<MDStoreRecord> source = mongoDatabase.getCollection(transactionId, MDStoreRecord.class);

		final MongoCollection<MDStoreRecord> destination = mongoDatabase.getCollection(currentId, MDStoreRecord.class)

				.withWriteConcern(WriteConcern.JOURNALED);

		//TODO reimplement using bulk writes on the destination collection

		for (MDStoreRecord record : source.find().batchSize(100).noCursorTimeout(true)) {

			final MDStoreRecord r = record.clone();

			if (StringUtils.isNotBlank(r.getId())) {

				destination.replaceOne(eq(ID, r.getId()), r, new UpdateOptions().upsert(true));

			}

		}

	}

	private MDStore getMDStore(final String mdId, final String currentId) {

		return new MongoMDStore(

				mdId,

				getRecordCollection(currentId),

				getRecordCollection(DISCARDED_PREFIX + currentId),

				metadata);

	}

	private MDStore getMDStore(final String id) {

		final Metadata m = metadata.findByMdId(id);

		return getMDStore(id, m.getCurrentId());

	}

	private void doDrop(final String mdId) {

		mongoDatabase.getCollection(mdId).drop();

		mongoDatabase.getCollection(DISCARDED_PREFIX + mdId).drop();

		log.debug("deleted collection " + mdId);

	}

	private MongoCollection<MDStoreRecord> getRecordCollection(final String collectionId) {

		return mongoDatabase.getCollection(collectionId, MDStoreRecord.class);

	}

	private void createIndex(final String fieldName, final MongoCollection<?> collection) {

		Arrays.asList(collection).forEach(c -> {

			log.info(String.format("Create index in %s", c.getNamespace().getCollectionName()));

			c.createIndex(

					new BasicDBObject(fieldName, 1),

					new IndexOptions().background(true));

		});

	}

	public int getMaxTransactions() {

		return maxTransactions;

	}

	public void setMaxTransactions(final int maxTransactions) {

		this.maxTransactions = maxTransactions;

	}

	public int getExpiredDays() {

		return expiredDays;

	}

	public void setExpiredDays(final int expiredDays) {

		this.expiredDays = expiredDays;

	}

	public MDStoreProfileManager getProfileManager() {

		return profileManager;

	}

	public void setProfileManager(final MDStoreProfileManager profileManager) {

		this.profileManager = profileManager;

	}

}