D-Net

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

import java.util.ArrayList;

import java.util.List;

import java.util.regex.Pattern;

import com.google.common.base.Function;

import com.google.common.collect.Iterables;

import com.google.common.collect.Lists;

import com.mongodb.BasicDBObject;

import com.mongodb.DBObject;

import com.mongodb.MongoClient;

import com.mongodb.client.FindIterable;

import com.mongodb.client.MongoCollection;

import com.mongodb.client.model.Sorts;

import eu.dnetlib.enabling.resultset.ResultSet;

import eu.dnetlib.enabling.resultset.ResultSetAware;

import eu.dnetlib.enabling.resultset.ResultSetListener;

import eu.dnetlib.miscutils.maps.ConcurrentSizedMap;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.bson.BsonDocument;

import org.bson.conversions.Bson;

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

public class MongoResultSetListener implements ResultSetListener, ResultSetAware {

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

	private ConcurrentSizedMap<Integer, String> lastKeys = new ConcurrentSizedMap<>();

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

	private Function<DBObject, String> serializer;

	private MongoCollection<DBObject> collection;

	private Bson query;

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

		this.collection = collection;

		this.serializer = serializer;

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

		log.debug("Query on mongo: "+this.query.toBsonDocument(BsonDocument.class, MongoClient.getDefaultCodecRegistry()));

	}

	@Override

	public List<String> getResult(final int fromPosition, final int toPosition) {

		ArrayList<DBObject> page = null;

		String lastKey = lastKeys.get(fromPosition);

		if (lastKey != null) {

			page = continueFrom(lastKey, (toPosition - fromPosition) + 1);

		} else {

			page = fetchNew(fromPosition - 1, (toPosition - fromPosition) + 1);

		}

		if (!page.isEmpty()) {

			DBObject last = page.get(page.size() - 1);

			lastKeys.put(toPosition + 1, (String) last.get("id"));

		}

		if (log.isDebugEnabled()) {

			log.info(String.format("got %s records from %s to %s", page.size(), fromPosition, toPosition));

		}

		return Lists.newArrayList(Iterables.transform(page, serializer));

	}

	private ArrayList<DBObject> fetchNew(final int from, final int size) {

		final FindIterable<DBObject> it = collection.find(query).batchSize(size);

		return Lists.newArrayList(it.sort(sortByIdAsc).skip(from).limit(size));

	}

	private ArrayList<DBObject> continueFrom(final String lastKey, final int size) {

		if (log.isDebugEnabled()) {

			log.debug("trying to continue from previous key: " + lastKey);

		}

		final Bson q = and(query, gt("id", lastKey));

		final FindIterable<DBObject> it = collection.find(q).batchSize(size).sort(sortByIdAsc).limit(size);

		return Lists.newArrayList(it);

	}

	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 int getSize() {

		return (int) collection.count(query);

	}

	@Override

	public void setResultSet(final ResultSet resultSet) {

		resultSet.close();

	}

}