D-Net

package eu.dnetlib.functionality.index.client.solr;

import java.io.IOException;

import java.text.SimpleDateFormat;

import java.util.*;

import com.google.common.collect.BiMap;

import com.google.common.collect.Maps;

import eu.dnetlib.data.provision.index.rmi.BrowsingRow;

import eu.dnetlib.data.provision.index.rmi.GroupResult;

import eu.dnetlib.data.provision.index.rmi.IndexServiceException;

import eu.dnetlib.functionality.cql.CqlValueTransformerMap;

import eu.dnetlib.functionality.index.client.IndexClient;

import eu.dnetlib.functionality.index.client.IndexClientException;

import eu.dnetlib.functionality.index.client.response.BrowseEntry;

import eu.dnetlib.functionality.index.client.response.BrowseValueEntry;

import eu.dnetlib.functionality.index.client.response.LookupResponse;

import eu.dnetlib.functionality.index.model.Any.ValueType;

import eu.dnetlib.functionality.index.query.*;

import eu.dnetlib.functionality.index.solr.cql.SolrTypeBasedCqlValueTransformerMapFactory;

import eu.dnetlib.functionality.index.solr.feed.StreamingInputDocumentFactory;

import eu.dnetlib.functionality.index.utils.MetadataReference;

import eu.dnetlib.functionality.index.utils.ZkServers;

import eu.dnetlib.miscutils.datetime.HumanTime;

import eu.dnetlib.miscutils.functional.UnaryFunction;

import org.apache.commons.lang3.StringUtils;

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.apache.solr.client.solrj.SolrClient;

import org.apache.solr.client.solrj.SolrQuery;

import org.apache.solr.client.solrj.SolrServerException;

import org.apache.solr.client.solrj.impl.CloudSolrClient;

import org.apache.solr.client.solrj.request.LukeRequest;

import org.apache.solr.client.solrj.response.LukeResponse;

import org.apache.solr.client.solrj.response.LukeResponse.FieldInfo;

import org.apache.solr.client.solrj.response.QueryResponse;

import org.apache.solr.client.solrj.response.UpdateResponse;

import org.apache.solr.common.SolrInputDocument;

/**

 * The Class SolrIndexClient.

 */

public class SolrIndexClient implements IndexClient {

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

	private static final String INDEX_RECORD_RESULT_FIELD = "dnetResult";

	private static String ZK_ADDRESS = "address";

	/** The format. */

	private String format;

	/** The layout. */

	private String layout;

	/** The interpretation. */

	private String interpretation;

	protected Map<String, String> serviceProperties;

	/** The client. */

	private CloudSolrClient client;

	private SolrIndexQueryFactory solrIndexQueryFactory;

	/** The query response factory. */

	private QueryResponseFactory<QueryResponse> queryResponseFactory;

	private SolrTypeBasedCqlValueTransformerMapFactory tMapFactory;

	/**

	 * The Constructor.

	 *

	 * @param format

	 *            the format

	 * @param layout

	 *            the layout

	 * @param interpretation

	 *            the interpretation

	 * @param serviceProperties

	 *            the service properties

	 * @param tMapFactory

	 */

	public SolrIndexClient(final String format, final String layout, final String interpretation, final Map<String, String> serviceProperties,

			final SolrIndexQueryFactory indexQueryFactory, final QueryResponseFactory<QueryResponse> queryResponseFactory,

			final SolrTypeBasedCqlValueTransformerMapFactory tMapFactory) {

		this.format = format;

		this.layout = layout;

		this.interpretation = interpretation;

		this.serviceProperties = serviceProperties;

		this.solrIndexQueryFactory = indexQueryFactory;

		this.queryResponseFactory = queryResponseFactory;

		this.tMapFactory = tMapFactory;

		log.debug(String.format("Created a new instance of the index of type %s-%s-%s", format, layout, interpretation));

	}

	/**

	 * Do delete.

	 *

	 * @param query

	 *            the CQL query

	 * @return true, if do delete

	 * @throws IndexServiceException

	 *             the index service exception

	 */

	@Override

	public long delete(final String query) throws IndexClientException {

		try {

			log.debug("delete by query: " + query);

			MetadataReference mdRef = new MetadataReference(getFormat(), getLayout(), getInterpretation());

			SolrIndexQuery translatedQuery = (SolrIndexQuery) solrIndexQueryFactory.getIndexQuery(QueryLanguage.CQL, query, this, mdRef);

			String tquery = translatedQuery.getQuery();

			translatedQuery.setQueryLimit(0);

			SolrIndexQueryResponse rsp = new SolrIndexQueryResponse(client.query(translatedQuery));

			QueryResponseParser responseParser = queryResponseFactory.getQueryResponseParser(rsp, mdRef);

			long total = responseParser.getNumFound();

			client.deleteByQuery(tquery);

			client.commit();

			return total;

		} catch (Exception e) {

			throw new IndexClientException("unable to run delete by query: " + query, e);

		}

	}

	/**

	 * {@inheritDoc}

	 *

	 * @throws IndexClientException

	 *

	 * @see IndexClient#browse(String, List, int)

	 */

	@Override

	public List<BrowseEntry> browse(final String query, final List<String> browseFields, final int max) throws IndexClientException {

		MetadataReference mdRef = new MetadataReference(getFormat(), getLayout(), getInterpretation());

		SolrIndexQuery translatedQuery = buildBrowseQuery(query, browseFields, max, mdRef);

		return executeBrowseQuery(query, translatedQuery, mdRef, browseFields);

	}

	/**

	 * {@inheritDoc}

	 *

	 * @throws IndexClientException

	 *

	 * @see IndexClient#browse(String, List, int, List)

	 */

	@Override

	public List<BrowseEntry> browse(final String query, final List<String> browseFields, final int max, final List<String> filterQuery)

			throws IndexClientException {

		MetadataReference mdRef = new MetadataReference(getFormat(), getLayout(), getInterpretation());

		SolrIndexQuery translatedQuery = buildBrowseQuery(query, browseFields, max, mdRef);

		if (filterQuery != null) {

			log.debug("Filter Query:");

			for (String fq : filterQuery) {

				translatedQuery.addFilterQuery(fq);

				log.debug("- " + fq);

			}

		}

		return executeBrowseQuery(query, translatedQuery, mdRef, browseFields);

	}

	private SolrIndexQuery buildBrowseQuery(final String query, final List<String> browseFields, final int max, final MetadataReference mdRef)

			throws IndexClientException {

		log.debug("Browse request for the index collection for query:" + query);

		SolrIndexQuery translatedQuery = (SolrIndexQuery) solrIndexQueryFactory.getIndexQuery(QueryLanguage.CQL, query, this, mdRef);

		translatedQuery.setFacet(true);

		if (browseFields != null) {

			List<String> browsableFields = solrIndexQueryFactory.getBrowsableFields(browseFields, mdRef);

			log.debug("Browsing fields:");

			for (String field : browsableFields) {

				translatedQuery.addFacetField(field);

				log.debug("- " + field);

			}

			translatedQuery.setFacetLimit(max);

			log.debug("max number of browsing field :" + max);

		}

		return translatedQuery;

	}

	private List<BrowseEntry> executeBrowseQuery(final String originalQuery,

			final SolrIndexQuery query,

			final MetadataReference mdRef,

			final List<String> browseFields) throws IndexClientException {

		try {

			SolrIndexQueryResponse response = new SolrIndexQueryResponse(client.query(query));

			QueryResponseParser responseParser = queryResponseFactory.getQueryResponseParser(response, mdRef);

			List<BrowsingRow> results = responseParser.getBrowsingResults();

			List<BrowseEntry> out = convertBrowseEntry(browseFields, results, responseParser.getAliases());

			return out;

		} catch (SolrServerException | IOException e) {

			throw new IndexClientException("Error on executing a query " + originalQuery, e);

		}

	}

	/**

	 * Creates the connection.

	 *

	 * @return the string

	 * @throws IndexClientException

	 *             the index client exception

	 */

	private String getUrl() throws IndexClientException {

		String address = serviceProperties.get(ZK_ADDRESS);

		if (StringUtils.isBlank(address)) {

			throw new IndexClientException("Unable to load a solr client, missing zk address");

		}

		return address;

	}

	/**

	 * Gets the client.

	 *

	 * @return the client

	 * @throws IndexClientException

	 *             the index client exception

	 */

	public SolrClient getClient() throws IndexClientException {

		if (this.client == null) {

			String url = getUrl();

			log.debug("create new Client " + url);

			final ZkServers zk = ZkServers.newInstance(url);

			final CloudSolrClient client = new CloudSolrClient.Builder(zk.getHosts(), zk.getChroot()).build();

			client.connect();

			client.setDefaultCollection(String.format("%s-%s-%s", getFormat(), getLayout(), getInterpretation()));

			try {

				client.ping();

			} catch (Exception e) {

				throw new IndexClientException("oops something went wrong", e);

			}

		}

		return client;

	}

	/**

	 * Sets the client.

	 *

	 * @param client

	 *            the client

	 */

	public void setClient(final CloudSolrClient client) {

		this.client = client;

	}

	@Override

	public LookupResponse lookup(final String query, final List<String> filterQuery, final int from, final int to) throws IndexClientException {

		log.debug("lookup request for the index collection for query:" + query);

		MetadataReference mdRef = new MetadataReference(getFormat(), getLayout(), getInterpretation());

		SolrIndexQuery translatedQuery = (SolrIndexQuery) solrIndexQueryFactory.getIndexQuery(QueryLanguage.CQL, query, this, mdRef);

		translatedQuery.setQueryOffset(from);

		translatedQuery.setQueryLimit(to - from + 1);

		if (filterQuery != null) {

			for (String fq : filterQuery) {

				translatedQuery.addFilterQuery(fq);

			}

		}

		try {

			SolrIndexQueryResponse response = new SolrIndexQueryResponse(client.query(translatedQuery));

			QueryResponseParser responseParser = queryResponseFactory.getQueryResponseParser(response, mdRef);

			return new LookupResponse(responseParser);

		} catch (SolrServerException | IOException e) {

			throw new IndexClientException("Error on executing a query " + query, e);

		}

	}

	@Override

	public CqlValueTransformerMap getCqlValueTransformerMap(final MetadataReference mdRef) throws IndexClientException {

		try {

			return tMapFactory.getIt(readFieldNamesAndTypes());

		} catch (Exception e) {

			throw new IndexClientException(e);

		}

	}

	@Override

	public IndexQueryFactory getIndexQueryFactory() {

		return solrIndexQueryFactory;

	}

	@Override

	public void close() throws IOException {

		log.debug("shutdown client: " + serviceProperties.get(ZK_ADDRESS));

		client.close();

	}

	public int feed(final String record, final String indexDsId, final UnaryFunction<String, String> toIndexRecord) throws IndexClientException {

		return feed(record, indexDsId, toIndexRecord, true);

	}

	public int feed(final String record, final String indexDsId, final UnaryFunction<String, String> toIndexRecord, final boolean commit)

			throws IndexClientException {

		try {

			final SolrInputDocument doc = prepareSolrDocument(record, indexDsId, toIndexRecord);

			if ((doc == null) || doc.isEmpty()) throw new IndexClientException("Invalid solr document");

			return feed(doc, commit);

		} catch (final Throwable e) {

			throw new IndexClientException("Error feeding document", e);

		}

	}

	public int feed(final SolrInputDocument document) throws IndexClientException {

		return feed(document, true);

	}

	public int feed(final List<SolrInputDocument> document) throws IndexClientException {

		try {

			final UpdateResponse res = client.add(document);

			log.debug("feed time for single records, elapsed time: " + HumanTime.exactly(res.getElapsedTime()));

			if (res.getStatus() != 0) { throw new IndexClientException("bad status: " + res.getStatus()); }

			return res.getStatus();

		} catch (final Throwable e) {

			throw new IndexClientException("Error feeding document", e);

		}

	}

	public int feed(final SolrInputDocument document, final boolean commit) throws IndexClientException {

		try {

			final UpdateResponse res = client.add(document);

			log.debug("feed time for single records, elapsed time: " + HumanTime.exactly(res.getElapsedTime()));

			if (res.getStatus() != 0) { throw new IndexClientException("bad status: " + res.getStatus()); }

			if (commit) {

				client.commit();

			}

			return res.getStatus();

		} catch (final Throwable e) {

			throw new IndexClientException("Error feeding document", e);

		}

	}

	public void feed(final List<SolrInputDocument> docs, final AfterFeedingCallback callback) throws IndexClientException {

		feed(docs, callback, true);

	}

	public void feed(final List<SolrInputDocument> docs, final AfterFeedingCallback callback, final boolean commit) throws IndexClientException {

		try {

			if (docs.isEmpty()) {

				log.debug("Empty list of documents. Calling callback, if needed.");

				if (callback != null) {

					callback.doAfterFeeding(null);

				}

				return;

			}

			final UpdateResponse res = client.add(docs);

			log.debug("feed time for " + docs.size() + " records, elapsed tipe: : " + HumanTime.exactly(res.getElapsedTime()));

			if (commit) {

				client.commit();

			}

			if (callback != null) {

				callback.doAfterFeeding(res);

			}

			if (res.getStatus() != 0) throw new IndexClientException("bad status: " + res.getStatus());

		} catch (final Throwable e) {

			throw new IndexClientException("Error feeding documents", e);

		}

	}

	public SolrInputDocument prepareSolrDocument(final String record, final String indexDsId, final UnaryFunction<String, String> toIndexRecord)

			throws IndexClientException {

		try {

			final StreamingInputDocumentFactory documentFactory = new StreamingInputDocumentFactory();

			final String version = (new SimpleDateFormat("yyyy-MM-dd\'T\'hh:mm:ss\'Z\'")).format(new Date());

			final String indexRecord = toIndexRecord.evaluate(record);

			if (log.isDebugEnabled()) {

				log.debug("***************************************\nSubmitting index record:\n" + indexRecord + "\n***************************************\n");

			}

			return documentFactory.parseDocument(version, indexRecord, indexDsId, INDEX_RECORD_RESULT_FIELD);

		} catch (final Throwable e) {

			throw new IndexClientException("Error creating solr document", e);

		}

	}

	public boolean isRecordIndexed(final String id) throws IndexClientException {

		final QueryResponse res = query("objidentifier:\"" + id + "\"", null);

		return res.getResults().size() > 0;

	}

	public int remove(final String id) throws IndexClientException {

		return remove(id, true);

	}

	public UpdateResponse commit() throws IndexClientException {

		try {

			return client.commit();

		} catch (SolrServerException | IOException e) {

			throw new IndexClientException(e);

		}

	}

	public int remove(final String id, final boolean commit) throws IndexClientException {

		try {

			final UpdateResponse res = client.deleteByQuery("objidentifier:\"" + id + "\"");

			if (commit) {

				client.commit();

			}

			return res.getResponse().size();

		} catch (final Throwable e) {

			throw new IndexClientException("Error removing documents", e);

		}

	}

	public int count(final String query) throws IndexClientException {

		final QueryResponse res = query(query, 0);

		return res.getResults().size();

	}

	public QueryResponse query(final String query, Integer rows) throws IndexClientException {

		try {

			final SolrQuery solrQuery = new SolrQuery();

			solrQuery.setQuery(query);

			if(rows != null && rows >= 0) {

				solrQuery.setRows(rows);

			}

			return client.query(solrQuery);

		} catch (final Throwable e) {

			throw new IndexClientException("Error searching documents", e);

		}

	}

	public UpdateResponse deleteByQuery(final String query) throws IndexClientException {

		try {

			return client.deleteByQuery(query);

		} catch (final Throwable e) {

			throw new IndexClientException("Error searching documents", e);

		}

	}

	public interface AfterFeedingCallback {

		void doAfterFeeding(final UpdateResponse response);

	}

	/**

	 * Gets the format.

	 *

	 * @return the format

	 */

	public String getFormat() {

		return format;

	}

	/**

	 * Sets the format.

	 *

	 * @param format

	 *            the format

	 */

	public void setFormat(final String format) {

		this.format = format;

	}

	/**

	 * Gets the layout.

	 *

	 * @return the layout

	 */

	public String getLayout() {

		return layout;

	}

	/**

	 * Sets the layout.

	 *

	 * @param layout

	 *            the layout

	 */

	public void setLayout(final String layout) {

		this.layout = layout;

	}

	/**

	 * Gets the interpretation.

	 *

	 * @return the interpretation

	 */

	public String getInterpretation() {

		return interpretation;

	}

	/**

	 * Sets the interpretation.

	 *

	 * @param interpretation

	 *            the interpretation

	 */

	public void setInterpretation(final String interpretation) {

		this.interpretation = interpretation;

	}

	public Map<String, String> getServiceProperties() {

		return serviceProperties;

	}

	public void setServiceProperties(final Map<String, String> serviceProperties) {

		this.serviceProperties = serviceProperties;

	}

	// HELPERS

	private List<BrowseEntry> convertBrowseEntry(final List<String> browseFields, final List<BrowsingRow> results, final BiMap<String, String> aliases) {

		Map<String, BrowseEntry> mapResult = new HashMap<String, BrowseEntry>();

		for (BrowsingRow row : results) {

			for (GroupResult groupResult : row.getGroupResult()) {

				String name = groupResult.getName();

				List<BrowseValueEntry> valuesEntry;

				BrowseEntry entry;

				if (mapResult.containsKey(name)) {

					entry = mapResult.get(name);

					valuesEntry = entry.getValues();

					if (valuesEntry == null) {

						valuesEntry = new ArrayList<BrowseValueEntry>();

						entry.setValues(valuesEntry);

					}

				} else {

					entry = new BrowseEntry();

					entry.setField(name);

					entry.setLabel(name);

					valuesEntry = new ArrayList<BrowseValueEntry>();

					entry.setValues(valuesEntry);

					mapResult.put(name, entry);

				}

				String value = groupResult.getValue();

				int count = groupResult.getCount();

				BrowseValueEntry entryValue = new BrowseValueEntry(value, count);

				valuesEntry.add(entryValue);

			}

		}

		List<BrowseEntry> out = new ArrayList<BrowseEntry>();

		for (String b : browseFields) {

			String inverse = null;

			if (aliases != null) {

				inverse = aliases.get(b) != null ? aliases.get(b) : aliases.inverse().get(b);

			}

			if (mapResult.containsKey(b)) {

				out.add(mapResult.get(b));

			} else if (mapResult.containsKey(inverse) == true) {

				BrowseEntry data = mapResult.get(inverse);

				data.setField(b);

				out.add(data);

			}

		}

		return out;

	}

	private Map<String, ValueType> readFieldNamesAndTypes() throws SolrServerException, IOException, IndexClientException {

		final LukeRequest request = new LukeRequest();

		request.setShowSchema(true);

		request.setNumTerms(0);

		final LukeResponse response = request.process(getClient());

		final Map<String, FieldInfo> fieldInfos = response.getFieldInfo();

		final Map<String, LukeResponse.FieldTypeInfo> fieldTypeInfos = response.getFieldTypeInfo();

		final Map<String, ValueType> result = Maps.newHashMap();

		for (FieldInfo fieldInfo : fieldInfos.values()) {

			LukeResponse.FieldTypeInfo fieldTypeInfo = fieldTypeInfos.get(fieldInfo.getType());

			final String fieldName = fieldTypeInfo.getName().toLowerCase();

			final ValueType fieldType = resolveSolrTypeClassName(fieldName);

			result.put(fieldInfo.getName(), fieldType);

		}

		return result;

	}

	private ValueType resolveSolrTypeClassName(final String solrTypeName) {

		if (solrTypeName.contains("LongField")) {

			return ValueType.LONG;

		} else if (solrTypeName.contains("IntField")) {

			return ValueType.LONG;

		} else if (solrTypeName.contains("short")) {

			return ValueType.LONG;

		} else if (solrTypeName.contains("float")) {

			return ValueType.DOUBLE;

		} else if (solrTypeName.contains("double")) {

			return ValueType.DOUBLE;

		} else if (solrTypeName.contains("date")) {

			return ValueType.DATETIME;

		} else {

			return ValueType.STRING;

		}

	}

}