org.eclipse.opencert.elastic/src/org/eclipse/opencert/elastic/search/ElasticFinderImpl.java

/*******************************************************************************
 * Copyright (C) 2018 ANSYS medini Technologies AG
 * 
 * This program and the accompanying materials are made
 * available under the terms of the Eclipse Public License 2.0
 * which is available at https://www.eclipse.org/legal/epl-2.0/
 * 
 * SPDX-License-Identifier: EPL-2.0
 * 
 * Contributors: 
 * 	ANSYS medini Technologies AG - initial API and implementation
 ******************************************************************************/
package org.eclipse.opencert.elastic.search;

import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.apache.http.HttpHost;
import org.eclipse.opencert.elastic.ElasticClient;
import org.eclipse.opencert.elastic.ElasticDocument;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.query.QueryStringQueryBuilder;
import org.elasticsearch.index.query.TermQueryBuilder;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;

import com.google.gson.JsonObject;

/**
 * Factory for the {@link ElasticFinder}.
 *
 * @author mauersberger
 */
public class ElasticFinderImpl implements ElasticFinder {

	// XXX hack to get some API tests running
	private Collection<Object> dummyData = null;

	private RestClient client;
	
	private String index = "_all";
	
	@SuppressWarnings("unused")
	private int lastStatus;
	
	/**
	 * Factory method to create a new {@link ElasticFinder}. XXX This is an
	 * intermediate feature.
	 * 
	 * @param data
	 * @return
	 */
	public static ElasticFinder onDummy(Object data) {
		ElasticFinderImpl impl = new ElasticFinderImpl();
		impl.dummyData = DummyData.INSTANCE.convert(data);
		return impl;
	}

	public static ElasticFinder onDefaultClient() {
		RestClient client = RestClient.builder(new HttpHost("localhost", 9200, "http")).build();
		return onClient(client);
	}

	public static ElasticFinder onClient(RestClient client) {
		ElasticFinderImpl impl = new ElasticFinderImpl();
		impl.client = client;
		return impl;
	}

	public static ElasticFinder onClient(ElasticClient client) {
		ElasticFinderImpl impl = new ElasticFinderImpl();
		impl.client = RestClient.builder(client.endPoint()).build();;
		return impl;
	}

	/*
	 * Intentionally private.
	 */
	private ElasticFinderImpl() {
		// just to avoid instantiation, force using the factory API
	}

	@Override
	public ElasticFinder within(String index) {
		this.index = (index != null ? index : "_all");
		return this;
	}
	
	@Override
	public Set<Hit> search(String query, Map<String, Object> filters) {
		if (dummyData != null) {
			return searchInDummyData(query, filters);
		}
		
		try {
			Set<Hit> hits = new HashSet<>();
			Set<ElasticDocument> documents = this.searchInElastic(index, query, filters);
			for (ElasticDocument document : documents) {
				Hit hit = new Hit();
				// we count 1 for string match + number of filters
				hit.score = 1 + (filters != null ? filters.size() : 0);
				hit.objectId = document.id;
				hits.add(hit);
			}
			return hits;
		} catch (Exception e) {
			e.printStackTrace();
			return Collections.emptySet();
		}
	}

	private Set<ElasticDocument> searchInElastic(String indexName, String queryString, Map<String, Object> filters)  throws Exception {
		RestHighLevelClient highLevel = new RestHighLevelClient(this.client);
		SearchRequest searchRequest = new SearchRequest(indexName);
		// see example at https://dzone.com/articles/java-high-level-rest-client-elasticsearch
		SearchSourceBuilder sourceBuilder = new SearchSourceBuilder();
		searchRequest.source(sourceBuilder);
		// find the best query builder depending on the arguments
		QueryBuilder queryBuilder = QueryBuilders.matchAllQuery();
		if (queryString != null) {
			queryString = queryString.replaceAll("\\.\\*", "*");
			QueryStringQueryBuilder stringQuery = QueryBuilders.queryStringQuery(queryString);
			stringQuery.allowLeadingWildcard(true);
			queryBuilder = stringQuery;
		}
		
		// for each filter we add a separate query that MUST match
		if (filters != null && filters.size() > 0) {
			// add the original query
			BoolQueryBuilder boolQuery = QueryBuilders.boolQuery();
			boolQuery.must(queryBuilder);
			
			// now create and add a quers for each field
			Set<String> keys = filters.keySet();
			for (String key : keys) {
				Object value = filters.get(key);
				TermQueryBuilder termQuery = QueryBuilders.termQuery(key, value);
				boolQuery.must(termQuery);
			}
			assert boolQuery.must().size() == 1 + filters.size();
			queryBuilder = boolQuery;
		}
		
		sourceBuilder.query(queryBuilder);
		// each SearchHit be returned with an explanation of the hit (ranking)
		sourceBuilder.explain(true);
		
		SearchResponse response = highLevel.search(searchRequest);
		lastStatus = response.status().getStatus();

		Set<ElasticDocument> result = new HashSet<>();
		SearchHits hits = response.getHits();
		// Note: total hits might be much bigger 
		for (SearchHit hit : hits.getHits()) {
			// convert Hit to Json
			JsonObject json = new JsonObject();
			
			// put all source information back into JSON
			Map<String, Object> fields = hit.getSource();
			Set<String> keys = fields.keySet();
			for (String key : keys) {
				Object field = fields.get(key);
				json.addProperty(key, (String) field.toString()); 
			}
			
			ElasticDocument doc = new ElasticDocument(json, hit.getIndex(), hit.getType(), hit.getId());
			result.add(doc);
		}
		
		return result;
	}
	
	/*
	 * XXX Hack
	 */
	private Set<Hit> searchInDummyData(String query, Map<String, Object> filters) {
		Set<Hit> hits = new HashSet<>();
		for (Object data : dummyData) {
			String document = DummyData.INSTANCE.asDocument(data);
			if ((document.contains(query) || document.matches(query)) && matchFilters(document, filters)) {
				Hit hit = new Hit();
				// we count 1 for string match + number of filters
				hit.score = 1 + (filters != null ? filters.size() : 0);
				hit.objectId = DummyData.INSTANCE.getId(document);
				hits.add(hit);
			}
		}
		return hits;
	}

	/*
	 * XXX Hack
	 */
	private boolean matchFilters(String document, Map<String, Object> filters) {
		// check filters
		if (filters == null || filters.isEmpty()) {
			return true;
		}

		for (String name : filters.keySet()) {
			String filter = name + ": " + filters.get(name).toString();
			if (!document.contains(filter)) {
				return false;
			}
		}
		// all filters passed
		return true;
	}
}