D-Net

package eu.dnetlib.resolver.mdstore.plugin;

import com.mongodb.BasicDBObject;

import com.mongodb.MongoNamespace;

import com.mongodb.client.ListIndexesIterable;

import com.mongodb.client.MongoCollection;

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

import eu.dnetlib.data.mdstore.modular.action.DoneCallback;

import eu.dnetlib.data.mdstore.modular.action.FailedCallback;

import eu.dnetlib.data.mdstore.modular.action.FeedAction;

import eu.dnetlib.data.mdstore.modular.action.MDStorePlugin;

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

import eu.dnetlib.data.mdstore.modular.mongodb.MDStoreTransactionManagerImpl;

import eu.dnetlib.rmi.data.MDStoreServiceException;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.Document;

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

import java.io.BufferedReader;

import java.io.InputStreamReader;

import java.util.Map;

import static org.apache.commons.lang3.StringUtils.isNotBlank;

public class DLIMergeRecord implements MDStorePlugin {

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

    @Autowired

    private MDStoreTransactionManagerImpl transactionManager;

    @Autowired

    private MDStoreFeeder mdStoreFeeder;

    @Override

    public void run(MDStoreDao dao, Map<String, String> params, DoneCallback doneCallback, FailedCallback failedCallback) {

        final String id = params.get("mdStoreId");

        final String host = params.get("mongoHost");

        final String nsPrefix = params.get("nsPrefix");

        final String sparkPath = params.get("sparkPath");

        final String sparkJobPath = params.get("sparkJobPath");

        final String sparkApplicationName = params.get("sparkApplicationName");

        final String mongoDBName = params.get("mongoDBName");

        final String number_of_core = params.get("numExecutor");

        if (isNotBlank(id) && isNotBlank(host) && isNotBlank(nsPrefix) && isNotBlank(sparkJobPath) && isNotBlank(sparkPath)) {

            try {

                log.debug("starting spark job");

                final String mdStoreCollection = transactionManager.getMDStoreCollection(id);

                final String[] command = {sparkPath + "bin/spark-submit", sparkJobPath, host, transactionManager.getDb().getName(), mdStoreCollection, nsPrefix, number_of_core, sparkApplicationName};

                final ProcessBuilder builder = new ProcessBuilder(command);

                final Process p = builder.start();

                final BufferedReader reader = new BufferedReader(new InputStreamReader(p.getErrorStream()));

                String line;

                while ((line = reader.readLine()) != null) {

                    log.info(line);

                }

                p.waitFor();

                if (p.exitValue() != 0) {

                    throw new MDStoreServiceException("The spark job exit with error");

                }

                log.info("Merging complete... creating index in the new collection");

                ListIndexesIterable<Document> documents = transactionManager.getDb().getCollection(mdStoreCollection).listIndexes();

                MongoCollection<Document> outMdStore = transactionManager.getDb().getCollection("out" + mdStoreCollection);

                for (Document d : documents) {

                    d.get("key", Map.class).keySet().forEach(it -> {

                        String s = it.toString();

                        outMdStore.createIndex(new BasicDBObject(s, 1));

                    });

                }

                log.info("index Created, dropping old collection and rename the new one");

                transactionManager.getDb().getCollection(mdStoreCollection).drop();

                int size = (int) outMdStore.count();

                outMdStore.renameCollection(new MongoNamespace(mongoDBName, mdStoreCollection));

                mdStoreFeeder.touch(id, size);

                doneCallback.call(params);

            } catch (Throwable e) {

                throw new RuntimeException(e);

            }

        } else {

            throw new RuntimeException("missing one of the following parameters {mdStoreId,mongoHost,nsPrefix,sparkPath,sparkJobPath}");

        }

    }

    @Override

    public String getStatus() {

        return "30/100";

    }

}