D-Net

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

import java.util.function.Function;

import java.util.regex.Pattern;

import com.mongodb.BasicDBObject;

import com.mongodb.DBObject;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.MongoCursor;

import com.mongodb.client.model.Sorts;

import eu.dnetlib.enabling.resultset.listener.ResultSetListener;

import eu.dnetlib.rmi.common.ResultSetException;

import eu.dnetlib.rmi.data.MDStoreServiceException;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.conversions.Bson;

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

public class MongoResultSetListener implements ResultSetListener<String> {

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

	private MongoCursor<DBObject> currentCursor;

	private Bson sortByIdAsc = Sorts.orderBy(Sorts.ascending("id"));

	private Function<DBObject, String> serializer;

	private MongoCollection<DBObject> collection;

	private int total = 0;

	private int count = 0;

	private int batchSize;

	private Bson query;

	public MongoResultSetListener(final MongoCollection<DBObject> collection, final Long from, final Long until, final Pattern filter, final Function<DBObject, String> serializer, final int batchSize)

			throws MDStoreServiceException {

		this.collection = collection;

		this.serializer = serializer;

		this.batchSize = batchSize;

		this.query = query(from, until, filter);

		createCursor();

	}

	private void createCursor() throws MDStoreServiceException {

		log.debug("init mdstore cursor using query " + query);

		try {

			total = (int) collection.count(query);

//			currentCursor = collection.find(query).batchSize(batchSize).sort(sortByIdAsc).noCursorTimeout(true).iterator();

            currentCursor = collection.find(query).batchSize(batchSize).noCursorTimeout(true).iterator();

		} catch (Throwable e) {

			throw new MDStoreServiceException("Error on initialize the Mongodb cursor");

		}

	}

	private Bson query(final Long from, final Long until, final Pattern pattern) {

		final Bson dateFilter = dateQuery(from, until);

		final Bson regexFilter = regexQuery(pattern);

		if (dateFilter != null & regexFilter != null) {

			return and(dateFilter, regexFilter);

		} else if (dateFilter != null) {

			return dateFilter;

		} else if (regexFilter != null) {

			return regexFilter;

		}

		return new BasicDBObject();

	}

	private Bson dateQuery(final Long from, final Long until) {

		if (from != null & until != null) {

			return and(gt("timestamp", from), lt("timestamp", until));

		}

		if (from != null) {

			return gt("timestamp", from);

		}

		if (until != null) {

			return lt("timestamp", until);

		}

		return null;

	}

	private Bson regexQuery(final Pattern pattern) {

		if (pattern != null) {

			return regex("body", pattern);

		}

		return null;

	}

	@Override

	public boolean hasNext() {

		return currentCursor.hasNext();

	}

	@Override

	public String next() throws ResultSetException {

		count++;

		return serializer.apply(currentCursor.next());

	}

	@Override

	public int getCount() {

		return count;

	}

	@Override

	public int getTotal() {

		return total;

	}

}