D-Net

package eu.dnetlib.pid.resolver.store;

import com.mongodb.*;

import com.mongodb.client.FindIterable;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.MongoCursor;

import com.mongodb.client.MongoDatabase;

import com.mongodb.client.model.UpdateOptions;

import eu.dnetlib.enabling.tools.DnetStreamSupport;

import eu.dnetlib.miscutils.functional.xml.DnetXsltFunctions;

import eu.dnetlib.pid.resolver.model.ResolvedObject;

import eu.dnetlib.pid.resolver.model.factory.ResolvedObjectFactory;

import org.apache.commons.lang3.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.Document;

import org.bson.conversions.Bson;

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

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

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

import java.util.ArrayList;

import java.util.Arrays;

import java.util.List;

import java.util.stream.Collectors;

import java.util.stream.StreamSupport;

/**

 * Created by sandro on 9/7/16.

 */

public class ResolverStore {

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

	@Value("${services.pid.resolver.header}")

	private String ns_header;

	private static List<String> requirededIndex = new ArrayList<>(Arrays.asList(

			"pid",

			"pid_type",

			"resolver",

			"dnet_identifier",

			"typology"

	));

	private final UpdateOptions upsert = new UpdateOptions().upsert(true);

	private MongoClient mongoClient;

	private MongoDatabase mongoDatabase;

	private MongoCollection<DBObject> resolverCollection;

	private String collectionName;

	private String databaseName;

	@Autowired

	private ResolvedObjectFactory objectFactory;

	public void checkIntegrityCollection() {

		getMongoDatabase().getCollection(collectionName);

		final MongoCollection<DBObject> collection = getMongoDatabase().getCollection(collectionName, DBObject.class);

		List<String> idxName = StreamSupport.stream(collection.listIndexes().spliterator(), false)

				.map(document -> ((Document) document.get("key")).keySet().iterator().next())

				.collect(Collectors.toList());

		final List<String> toAdd = requirededIndex.stream()

				.filter(it -> !idxName.contains(it))

				.collect(Collectors.toList());

		for (String idx : toAdd) {

			log.debug("Adding " + idx + " index");

			collection.createIndex(new Document(idx, 1));

		}

		this.resolverCollection = collection;

	}

	public List<String> getIndexName() {

		return StreamSupport.stream(this.resolverCollection.listIndexes().spliterator(), false)

				.map(document -> ((Document) document.get("key")).keySet().iterator().next())

				.collect(Collectors.toList());

	}

	public void insertRecord(final String resolver, final ResolvedObject resolvedObject) {

		if (this.resolverCollection == null) {

			checkIntegrityCollection();

		}

		final DBObject currentObject = BasicDBObjectBuilder.start()

                .add("pid", resolvedObject.getPid().toLowerCase())

                .add("pid_type", resolvedObject.getPidType().toLowerCase())

                .add("dnet_identifier", generateDNetIdentifier(resolvedObject))

				.add("resolver", resolver)

				.add("typology", resolvedObject.getType().toString())

				.add("body", resolvedObject.toJsonString())

				.get();

		final Bson dnet_identifier = (Bson) BasicDBObjectBuilder.start().add("dnet_identifier", ns_header + resolvedObject.getIdentifier()).get();

		this.resolverCollection.replaceOne(dnet_identifier, currentObject, upsert);

	}

	public String generateDNetIdentifier(final String pid, final String pidType) {

		if (StringUtils.isBlank(pid) || StringUtils.isBlank(pidType))

			throw new RuntimeException("Error pid or pidtype cannot be null");

		return ns_header + DnetXsltFunctions.md5(String.format("%s::%s", pid.trim().toLowerCase(), pidType.toLowerCase().trim()));

	}

	public String generateDNetIdentifier(final ResolvedObject object) {

		return ns_header + object.getIdentifier();

	}

	public void insertRecords(final String resolver, final List<ResolvedObject> resolvedObjects) {

		resolvedObjects.forEach(it -> insertRecord(resolver, it));

	}

    public ResolvedObject getRecord(final String dnetIdentifier) {

        if (this.resolverCollection == null) {

            checkIntegrityCollection();

        }

        final Bson query = (Bson) QueryBuilder.start("dnet_identifier").is(dnetIdentifier).get();

        final FindIterable<DBObject> dbObjects = this.resolverCollection.find(query);

        if (dbObjects != null) {

            MongoCursor<DBObject> iterator = dbObjects.iterator();

			if (iterator.hasNext()) {

				final String body = iterator.next().get("body").toString();

				final ResolvedObject obj = objectFactory.generateObjectFromJson(body);

				return obj;

            }

        }

        return null;

    }

	public ResolvedObject getRecord(final String pid, final String pidType) {

		if (this.resolverCollection == null) {

			checkIntegrityCollection();

		}

        final Bson query = (Bson) QueryBuilder.start("pid").is(pid.toLowerCase()).and("pid_type").is(pidType.toLowerCase()).get();

        final FindIterable<DBObject> dbObjects = this.resolverCollection.find(query);

		if (dbObjects != null) {

			final List<ResolvedObject> rList = new ArrayList<>();

			dbObjects.forEach((Block<DBObject>) it -> {

				final String body = it.get("body").toString();

				ResolvedObject obj = objectFactory.generateObjectFromJson(it.get("body").toString());

				rList.add(obj);

			});

			if (rList.size() > 0) {

				return rList.get(0);

			}

		}

		return null;

	}

    public String getRecordIdentifier(final String dnetId) {

        double start = 0;

        try {

            if (this.resolverCollection == null) {

                checkIntegrityCollection();

            }

            start = System.currentTimeMillis();

            final Bson query = (Bson) QueryBuilder.start().put("dnet_identifier").is(dnetId).get();

            final FindIterable<DBObject> dbObjects = this.resolverCollection.find(query);

            if (dbObjects != null && dbObjects.iterator().hasNext()) {

                return dnetId;

            }

            return null;

        } finally {

            double total = System.currentTimeMillis() - start;

            log.debug("Request fetched in " + total + " ms");

        }

    }

	public String getRecordIdentifier(final String pid, final String pidType) {

		double start = 0;

		try {

			if (this.resolverCollection == null) {

				checkIntegrityCollection();

			}

			start = System.currentTimeMillis();

			final Bson query = (Bson) QueryBuilder.start().put("pid").is(pid).and("pid_type").is(pidType).get();

			final FindIterable<DBObject> dbObjects = this.resolverCollection.find(query);

			if (dbObjects != null) {

				for (DBObject obj : dbObjects) {

					final String id = obj.get("dnet_identifier").toString();

					if (id != null && !StringUtils.isBlank(id))

						return id;

				}

			}

			return null;

		} finally {

			double total = System.currentTimeMillis() - start;

			log.debug("Request fetched in " + total + " ms");

		}

	}

	public Iterable<ResolvedObject> getAllRecords() {

		try {

			if (this.resolverCollection == null) {

				checkIntegrityCollection();

			}

			final FindIterable<DBObject> dbObjects = this.resolverCollection.find();

			return () -> DnetStreamSupport.generateStreamFromIterator(dbObjects.iterator()).map(it ->

					objectFactory.generateObjectFromJson(it.get("body").toString())

			).iterator();

		} catch (Exception e) {

			return null;

		}

	}

	public MongoDatabase getMongoDatabase() {

		if (mongoDatabase == null) {

			mongoDatabase = mongoClient.getDatabase(getDatabaseName());

		}

		return mongoDatabase;

	}

	public void setMongoDatabase(final MongoDatabase mongoDatabase) {

		this.mongoDatabase = mongoDatabase;

	}

	public MongoCollection<DBObject> getResolverCollection() {

		return resolverCollection;

	}

	public void setResolverCollection(final MongoCollection<DBObject> resolverCollection) {

		this.resolverCollection = resolverCollection;

	}

	public String getCollectionName() {

		return collectionName;

	}

	@Required

	public void setCollectionName(final String collectionName) {

		this.collectionName = collectionName;

	}

	public String getDatabaseName() {

		return databaseName;

	}

	@Required

	public void setDatabaseName(final String databaseName) {

		this.databaseName = databaseName;

	}

	public MongoClient getMongoClient() {

		return mongoClient;

	}

	@Required

	public void setMongoClient(final MongoClient mongoClient) {

		this.mongoClient = mongoClient;

	}

}