D-Net

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

import java.util.List;

import com.google.common.collect.Lists;

import com.google.common.primitives.Ints;

import com.mongodb.BasicDBObject;

import com.mongodb.BasicDBObjectBuilder;

import com.mongodb.DBObject;

import com.mongodb.client.AggregateIterable;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.MongoDatabase;

import com.mongodb.client.model.CreateCollectionOptions;

import eu.dnetlib.data.mdstore.MDStoreServiceException;

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

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

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

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

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

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

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

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

import eu.dnetlib.miscutils.collections.FilteredCollection;

import eu.dnetlib.miscutils.collections.MappedCollection;

import eu.dnetlib.miscutils.datetime.DateUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.Document;

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

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

/**

 * Factory bean for MongoMDStore instances.

 *

 * @author marko

 */

public class MDStoreDaoImpl implements MDStoreDao {

	public static final String MD_ID = "mdId";

	public static final String FORMAT = "format";

	public static final String INTERPRETATION = "interpretation";

	public static final String LAYOUT = "layout";

	public static final String SIZE = "size";

	public static final String METADATA_NAME = "metadata";

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

	private MongoDatabase db;

	private RecordParserFactory recordParserFactory;

	private boolean discardRecords = true;

	@Autowired

	private MDStoreTransactionManager transactionManager;

	/**

	 * {@inheritDoc}

	 */

	@Override

	public void createMDStore(final String mdId, final String format, final String interpretation, final String layout) throws MDStoreServiceException {

		transactionManager.createMDStore(mdId);

		final String internalId = transactionManager.getMDStoreCollection(mdId);

		if (!Lists.newArrayList(getDb().listCollectionNames()).contains(internalId)) {

			log.info(String.format("creating collection %s", internalId));

			getDb().createCollection(internalId, new CreateCollectionOptions());

		}

		final MongoCollection<DBObject> coll = getDb().getCollection(METADATA_NAME, DBObject.class);

		final BasicDBObject obj = new BasicDBObject();

		obj.put(MD_ID, mdId);

		obj.put(FORMAT, format);

		obj.put(INTERPRETATION, interpretation);

		obj.put(LAYOUT, layout);

		obj.put(SIZE, 0);

		coll.insertOne(obj);

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public void deleteMDStore(final String mdId) throws MDStoreServiceException {

		final MongoCollection<DBObject> metadata = getDb().getCollection(METADATA_NAME, DBObject.class);

		if (metadata == null) throw new MDStoreServiceException("cannot find metadata collection");

		transactionManager.dropMDStore(mdId);

		metadata.deleteOne(new BasicDBObject(MD_ID, mdId));

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

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public MDStore getMDStore(final String mdId) throws MDStoreServiceException {

		final String internalId = transactionManager.getMDStoreCollection(mdId);

		return new MongoMDStore(mdId, getDb().getCollection(internalId, DBObject.class), getRecordParser(), isDiscardRecords(), getDb());

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public MDStore readMDStore(final String mdId) throws MDStoreServiceException {

		final String internalId = transactionManager.readMdStore(mdId);

		return new MongoMDStore(mdId, getDb().getCollection(internalId, DBObject.class), getRecordParser(), isDiscardRecords(), getDb());

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public MDStore startTransaction(final String mdId, final boolean refresh) throws MDStoreServiceException {

		final String transactionId = transactionManager.startTransaction(mdId, refresh);

		return new MongoMDStore(transactionId, getDb().getCollection(transactionId, DBObject.class), getRecordParser(), isDiscardRecords(),

				getDb());

	}

	private RecordParser getRecordParser() {

		final RecordParser parser = getRecordParserFactory().newInstance();

		parser.setTimestamp(DateUtils.now());

		return parser;

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public List<MDStoreDescription> listMDStores() {

		return MappedCollection.listMap(getDb().getCollection(METADATA_NAME, DBObject.class).find(), input -> {

			final String mdId = (String) input.get(MD_ID);

			log.debug("Getting info for " + mdId);

			final String format = (String) input.get(FORMAT);

			final String layout = (String) input.get(LAYOUT);

			final String interpretation = (String) input.get(INTERPRETATION);

			MongoMDStore currentMDStore = null;

			final MDStoreDescription description = new MDStoreDescription();

			try {

				currentMDStore = (MongoMDStore) getMDStore(mdId);

			} catch (final MDStoreServiceException e) {

				log.error("Error on retrieving mdstore for getting info mdId " + mdId);

			}

			int size = 0;

			if (input.containsField(SIZE)) {

				log.debug("Size retrieved from metadata for mdId :" + mdId);

				size = (Integer) input.get(SIZE);

			} else {

				if (currentMDStore != null) {

					log.debug("Size not Found in metadata for mdId :" + mdId + " calling getCount ");

					size = currentMDStore.getSize();

					input.put("size", size);

					getDb().getCollection(METADATA_NAME, DBObject.class).findOneAndReplace(new BasicDBObject(MD_ID, mdId), input);

				}

			}

			if (currentMDStore != null) {

				description.setIndexed(currentMDStore.isIndexed());

			}

			description.setId(mdId);

			description.setFormat(format);

			description.setLayout(layout);

			description.setInterpretation(interpretation);

			description.setSize(size);

			return description;

		});

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public List<String> listMDStores(final String format, final String layout, final String interpretation) {

		return MappedCollection.listMap(

				FilteredCollection.listFilter(getDb().getCollection(METADATA_NAME, DBObject.class).find(), MDStoreUtils.dboFilter(format, layout, interpretation)),

				MDStoreUtils.mdId());

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public int getCachedSize(final String id) throws MDStoreServiceException {

		log.debug("retrieve cached size for mdstore: " + id);

		final DBObject desc = getDb().getCollection(METADATA_NAME, DBObject.class).find(new BasicDBObject(MD_ID, id)).first();

		if (!desc.containsField(SIZE)) {

			desc.put(SIZE, getMDStore(id).getSize());

		}

		final Object oSize = desc.get(SIZE);

		return (Integer) oSize;

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public void refreshSizes() throws MDStoreServiceException {

		for (final MDStoreDescription mdStoreId : listMDStores()) {

			refreshSize(mdStoreId.getId());

		}

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public int refreshSize(final String mdStoreId) throws MDStoreServiceException {

		final int size = (int) getDb().getCollection(transactionManager.getMDStoreCollection(mdStoreId)).count();

		final MongoCollection<DBObject> metadata = getDb().getCollection(METADATA_NAME, DBObject.class);

		metadata.updateOne(new BasicDBObject(MD_ID, mdStoreId), new BasicDBObject("$set", new BasicDBObject(SIZE, size)));

		return size;

	}

	@Override

	public int getSumOfSizes(final String format, final String layout, final String interpretation) throws MDStoreServiceException {

		final MongoCollection<DBObject> metadata = getDb().getCollection(METADATA_NAME, DBObject.class);

		BasicDBObject matchObj = (BasicDBObject) BasicDBObjectBuilder.start("$match",

				BasicDBObjectBuilder.start("format", format).add("layout", layout).add("interpretation", interpretation).get()).get();

		BasicDBObject groupObj = (BasicDBObject) BasicDBObjectBuilder.start("$group",

				BasicDBObjectBuilder.start("_id", "").add("total", new BasicDBObject("$sum", "$" + SIZE)).get()).get();

		BasicDBObject projectObj = new BasicDBObject("$project", new BasicDBObject("_id", 0).append("total", 1));

		List<BasicDBObject> pipeline = Lists.newArrayList(matchObj, groupObj, projectObj);

		AggregateIterable<DBObject> output = metadata.aggregate(pipeline, DBObject.class);

		DBObject result = output.first();

		if (result == null || !result.containsField("total")) {

			log.debug("No total found");

			return 0;

		} else return (Integer) result.get("total");

	}

	/**

	 * {@inheritDoc}

	 */

	@Override

	public void commit(final String transactionId, final String mdId) throws MDStoreServiceException {

		transactionManager.commit(transactionId, mdId, getMDStore(mdId));

	}

	/**

	 * Getter for property 'db'.

	 *

	 * @return Value for property 'db'.

	 */

	public MongoDatabase getDb() {

		return db;

	}

	/**

	 * Setter for property 'db'.

	 *

	 * @param db

	 *            Value to set for property 'db'.

	 */

	@Required

	public void setDb(final MongoDatabase db) {

		this.db = db;

	}

	/**

	 * Getter for property 'recordParser'.

	 *

	 * @return Value for property 'recordParser'.

	 */

	public RecordParserFactory getRecordParserFactory() {

		return recordParserFactory;

	}

	@Required

	public void setRecordParserFactory(final RecordParserFactory recordParserFactory) {

		this.recordParserFactory = recordParserFactory;

	}

	/**

	 * Getter for property 'discardRecords'.

	 *

	 * @return Value for property 'discardRecords'.

	 */

	public boolean isDiscardRecords() {

		return discardRecords;

	}

	/**

	 * Setter for property 'discardRecords'.

	 *

	 * @param discardRecords

	 *            Value to set for property 'discardRecords'.

	 */

	public void setDiscardRecords(final boolean discardRecords) {

		this.discardRecords = discardRecords;

	}

	@Override

	public MDStoreDBStatus getDBStatus() {

		final int handledDatastructures = Ints.saturatedCast(getDb().getCollection(METADATA_NAME).count());

		//final int usedDiskSpace = Ints.saturatedCast(getDb().getStats().getLong("storageSize") / (1024 * 1024)); // in MB

		//{ dbStats: 1, scale: 1 }

		BasicDBObject statsQuery = new BasicDBObject("dbStats", 1);

		statsQuery.put("scale", 1024 * 1024); //storageSize in MB

		final Document statsRes = getDb().runCommand(statsQuery);

		log.debug("DBStatus --  " + statsRes.toJson());

		int usedDiskSpace;

		//trying to handle different versions of the mongo server: old version returns storage size as long, new version as double

		//TODO: simplify this when dev, beta, production are aligned with our local, latest, mongo version

		String usedDiskSpaceStr = statsRes.get("storageSize").toString();

		try {

			Long usedDiskSpaceLong = Long.parseLong(usedDiskSpaceStr);

			usedDiskSpace = Ints.saturatedCast(usedDiskSpaceLong);

		} catch (NumberFormatException nfe) {

			Double usedDiskSpaceDbl = Double.parseDouble(usedDiskSpaceStr);

			usedDiskSpace = usedDiskSpaceDbl.intValue();

		}

		final String date = DateUtils.now_ISO8601();

		return new MDStoreDBStatus(handledDatastructures, usedDiskSpace, date);

	}

	@Override

	public void startGarbage() throws MDStoreServiceException {

		this.transactionManager.garbage();

	}

	@Override

	public void invalidTransaction(final String transactionId, final String mdId) throws MDStoreServiceException {

		transactionManager.dropTransaction(mdId, transactionId);

	}

}