D-Net

package eu.dnetlib.functionality.index.utils;

import com.google.common.collect.Maps;

import com.google.gson.JsonElement;

import com.google.gson.JsonObject;

import com.google.gson.JsonParser;

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

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

import org.apache.commons.logging.Log;

import org.apache.commons.logging.LogFactory;

import org.apache.http.HttpResponse;

import org.apache.http.client.HttpClient;

import org.apache.http.client.methods.HttpGet;

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

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.SolrPingResponse;

import org.apache.solr.common.cloud.SolrZkClient;

import org.apache.zookeeper.WatchedEvent;

import org.apache.zookeeper.Watcher;

import org.apache.zookeeper.data.Stat;

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

import java.io.IOException;

import java.util.Map;

public class RemoteSolrAdministrator {

	/**

	 * The log.

	 */

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

	/** The create url request. */

	private static String createURLRequest = "http://%s:%s/solr/admin/collections?action=CREATE&name=%s&numShards=%s&replicationFactor=%s&maxShardsPerNode=%s&collection.configName=%s";

	/** The create url request. */

	private static String reloadURLRequest = "http://%s:%s/solr/admin/collections?action=RELOAD&name=%s";

	/** The http client. */

	private HttpClient httpClient;

	protected Map<String, Map<String, ValueType>> cachedSchema;

	public RemoteSolrAdministrator() {

		this.cachedSchema = Maps.newHashMap();

	}

	public boolean createSolrIndex(final String host,

			final String port,

			final String collectionName,

			final String numShard,

			final String replicationFactor,

			final String maxShardsPerNode,

		final String collectionConfigName) throws IndexServiceException {

		final String uri = generateCreateIndexRequest(host, port, collectionName, numShard, replicationFactor, maxShardsPerNode, collectionConfigName);

		log.info(uri);

		HttpGet request = new HttpGet(uri);

		HttpResponse response;

		try {

			response = getHttpClient().execute(request);

		} catch (Exception e) {

			throw new IndexServiceException("Unable to send request to solr server", e);

		} finally {

			request.releaseConnection();

		}

		if (response.getStatusLine().getStatusCode() != 200)

			throw new IndexServiceException("Error on creating index the error code from solr is" + response.toString());

		return false;

	}

	public boolean indexCollectionExists(final String indexCollectionId, final CloudSolrClient client) {

		Watcher watcher = new Watcher() {

			@Override

			public void process(final WatchedEvent arg0) {}

		};

		try {

			// server.connect();

			SolrZkClient zkClient = client.getZkStateReader().getZkClient();

			if (!zkClient.isConnected()) {

				client.close();

				client.connect();

			}

			byte[] data = zkClient.getData("/clusterstate.json", watcher, new Stat(), true);

			if (data == null) return false;

			String jsonLine = new String(data);

			JsonElement jelement = new JsonParser().parse(jsonLine);

			JsonObject jobject = jelement.getAsJsonObject();

			if (jobject.has(indexCollectionId)) {

				client.setDefaultCollection(indexCollectionId);

				client.connect();

				SolrPingResponse status = client.ping();

				return (Integer) status.getResponseHeader().get("status") == 0;

			} else return false;

		} catch (Exception e) {

			log.error(e);

			return false;

		}

	}

	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;

	}

	public Map<String, ValueType> getFieldNamesAndTypes(final String coreName, final SolrClient client) throws IndexServiceException {

		try {

			if (cachedSchema.get(coreName) == null) {

				synchronized (cachedSchema) {

					Map<String, ValueType> schema = readFieldNamesAndTypes(coreName, client);

					log.info("setting cache for schema of collection: " + coreName);

					cachedSchema.put(coreName, schema);

				}

			}

			return cachedSchema.get(coreName);

		} catch (Exception e) {

			throw new IndexServiceException("Unable to get Schema for " + coreName + " exception", e);

		}

	}

	private Map<String, ValueType> readFieldNamesAndTypes(final String coreName, final SolrClient client) throws SolrServerException, IOException {

		final LukeRequest request = new LukeRequest();

		request.setShowSchema(true);

		request.setNumTerms(0);

		final LukeResponse response = request.process(client);

		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 String generateCreateIndexRequest(final String host,

			final String port,

			final String collectionName,

			final String numShard,

			final String replicationFactor,

			final String collectionConfigName,

			final String maxShardsPerNode) {

		return String.format(createURLRequest, host, port, collectionName, numShard, replicationFactor, maxShardsPerNode, collectionConfigName);

	}

	private String generateUpdateIndexRequest(final String host, final String port, final String collectionName) {

		return String.format(reloadURLRequest, host, port, collectionName);

	}

	/**

	 * @return the httpClient

	 */

	public HttpClient getHttpClient() {

		return httpClient;

	}

	/**

	 * @param httpClient

	 *            the httpClient to set

	 */

	@Required

	public void setHttpClient(final HttpClient httpClient) {

		log.info("setting http client " + httpClient.getClass());

		this.httpClient = httpClient;

	}

	public void reloadCollection(final String host, final String port, final String collectionName) throws IndexServiceException {

		log.info("creating the request of reload index " + generateUpdateIndexRequest(host, port, collectionName));

		HttpGet request = new HttpGet(generateUpdateIndexRequest(host, port, collectionName));

		HttpResponse response;

		try {

			response = httpClient.execute(request);

		} catch (Exception e) {

			throw new IndexServiceException("Unable to send request to solr server", e);

		} finally {

			request.releaseConnection();

		}

		if (response.getStatusLine().getStatusCode() != 200)

			throw new IndexServiceException("Error on reloading index the error code from solr is" + response.toString());

	}

}