DATAES-631 - Consolidate query objects.

Author: sothawo

src/main/java/org/springframework/data/elasticsearch/core/AbstractElasticsearchTemplate.java

@@ -1,13 +1,43 @@
 package org.springframework.data.elasticsearch.core;
 
+import static org.springframework.util.StringUtils.*;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+import org.elasticsearch.action.bulk.BulkItemResponse;
+import org.elasticsearch.action.bulk.BulkResponse;
+import org.elasticsearch.action.search.MultiSearchRequest;
+import org.elasticsearch.action.search.MultiSearchResponse;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.common.collect.MapBuilder;
+import org.elasticsearch.index.query.MoreLikeThisQueryBuilder;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import org.springframework.beans.BeansException;
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.ApplicationContextAware;
+import org.springframework.data.domain.Page;
 import org.springframework.data.elasticsearch.ElasticsearchException;
+import org.springframework.data.elasticsearch.annotations.Document;
 import org.springframework.data.elasticsearch.annotations.Mapping;
+import org.springframework.data.elasticsearch.annotations.Setting;
 import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
 import org.springframework.data.elasticsearch.core.convert.MappingElasticsearchConverter;
+import org.springframework.data.elasticsearch.core.document.SearchDocumentResponse;
 import org.springframework.data.elasticsearch.core.index.MappingBuilder;
+import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
+import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
 import org.springframework.data.elasticsearch.core.mapping.SimpleElasticsearchMappingContext;
+import org.springframework.data.elasticsearch.core.query.DeleteQuery;
+import org.springframework.data.elasticsearch.core.query.MoreLikeThisQuery;
+import org.springframework.data.elasticsearch.core.query.NativeSearchQuery;
+import org.springframework.data.elasticsearch.core.query.NativeSearchQueryBuilder;
+import org.springframework.data.elasticsearch.core.query.Query;
+import org.springframework.util.Assert;
 import org.springframework.util.StringUtils;
 
 /**
@@ -16,11 +46,22 @@
  * @author Sascha Woo
  * @author Peter-Josef Meisch
  */
-public abstract class AbstractElasticsearchTemplate implements ElasticsearchOperations {
+public abstract class AbstractElasticsearchTemplate implements ElasticsearchOperations, ApplicationContextAware {
 
 	private static final Logger LOGGER = LoggerFactory.getLogger(AbstractElasticsearchTemplate.class);
 
 	protected ElasticsearchConverter elasticsearchConverter;
+	protected RequestFactory requestFactory;
+
+	public RequestFactory getRequestFactory() {
+		return requestFactory;
+	}
+
+	protected void initialize(ElasticsearchConverter elasticsearchConverter) {
+		Assert.notNull(elasticsearchConverter, "elasticsearchConverter must not be null.");
+		this.elasticsearchConverter = elasticsearchConverter;
+		this.requestFactory = new RequestFactory(elasticsearchConverter);
+	}
 
 	protected ElasticsearchConverter createElasticsearchConverter() {
 		MappingElasticsearchConverter mappingElasticsearchConverter = new MappingElasticsearchConverter(
@@ -29,6 +70,13 @@ protected ElasticsearchConverter createElasticsearchConverter() {
 		return mappingElasticsearchConverter;
 	}
 
+	@Override
+	public void setApplicationContext(ApplicationContext context) throws BeansException {
+		if (elasticsearchConverter instanceof ApplicationContextAware) {
+			((ApplicationContextAware) elasticsearchConverter).setApplicationContext(context);
+		}
+	}
+
 	protected String buildMapping(Class<?> clazz) {
 
 		// load mapping specified in Mapping annotation if present
@@ -53,8 +101,208 @@ protected String buildMapping(Class<?> clazz) {
 		}
 	}
 
+	@Override
+	public boolean createIndex(String indexName) {
+		return createIndexIfNotCreated(indexName);
+	}
+
+	private <T> boolean createIndexIfNotCreated(String indexName) {
+		return indexExists(indexName) || createIndex(indexName, null);
+	}
+
+	@Override
+	public <T> boolean createIndex(Class<T> clazz) {
+		return createIndexIfNotCreated(clazz);
+	}
+
+	private <T> boolean createIndexIfNotCreated(Class<T> clazz) {
+		return indexExists(getPersistentEntityFor(clazz).getIndexName()) || createIndexWithSettings(clazz);
+	}
+
+	private <T> boolean createIndexWithSettings(Class<T> clazz) {
+		if (clazz.isAnnotationPresent(Setting.class)) {
+			String settingPath = clazz.getAnnotation(Setting.class).settingPath();
+			if (hasText(settingPath)) {
+				String settings = ResourceUtil.readFileFromClasspath(settingPath);
+				if (hasText(settings)) {
+					return createIndex(getPersistentEntityFor(clazz).getIndexName(), settings);
+				}
+			} else {
+				LOGGER.info("settingPath in @Setting has to be defined. Using default instead.");
+			}
+		}
+		return createIndex(getPersistentEntityFor(clazz).getIndexName(), getDefaultSettings(getPersistentEntityFor(clazz)));
+	}
+
+	@Override
+	public <T> boolean createIndex(Class<T> clazz, Object settings) {
+		return createIndex(getPersistentEntityFor(clazz).getIndexName(), settings);
+	}
+
+	@Override
+	public void delete(Query query, Class<?> clazz, IndexCoordinates index) {
+		Assert.notNull(query, "Query must not be null.");
+		SearchRequest searchRequest = requestFactory.searchRequest(query, clazz, index);
+		DeleteQuery deleteQuery = new DeleteQuery();
+		deleteQuery.setQuery(searchRequest.source().query());
+		delete(deleteQuery, index);
+	}
+
+	@Override
+	public <T> Page<T> moreLikeThis(MoreLikeThisQuery query, Class<T> clazz, IndexCoordinates index) {
+		Assert.notNull(query.getId(), "No document id defined for MoreLikeThisQuery");
+		MoreLikeThisQueryBuilder moreLikeThisQueryBuilder = requestFactory.moreLikeThisQueryBuilder(query, index);
+		return queryForPage(new NativeSearchQueryBuilder().withQuery(moreLikeThisQueryBuilder).build(), clazz, index);
+	}
+
+	protected static String[] toArray(List<String> values) {
+		String[] valuesAsArray = new String[values.size()];
+		return values.toArray(valuesAsArray);
+	}
+
 	@Override
 	public ElasticsearchConverter getElasticsearchConverter() {
 		return elasticsearchConverter;
 	}
+
+	@Override
+	public ElasticsearchPersistentEntity getPersistentEntityFor(Class clazz) {
+		Assert.isTrue(clazz.isAnnotationPresent(Document.class), "Unable to identify index name. " + clazz.getSimpleName()
+				+ " is not a Document. Make sure the document class is annotated with @Document(indexName=\"foo\")");
+		return elasticsearchConverter.getMappingContext().getRequiredPersistentEntity(clazz);
+	}
+
+	private <T> Map getDefaultSettings(ElasticsearchPersistentEntity<T> persistentEntity) {
+
+		if (persistentEntity.isUseServerConfiguration())
+			return new HashMap();
+
+		return new MapBuilder<String, String>().put("index.number_of_shards", String.valueOf(persistentEntity.getShards()))
+				.put("index.number_of_replicas", String.valueOf(persistentEntity.getReplicas()))
+				.put("index.refresh_interval", persistentEntity.getRefreshInterval())
+				.put("index.store.type", persistentEntity.getIndexStoreType()).map();
+	}
+
+	protected void checkForBulkOperationFailure(BulkResponse bulkResponse) {
+		if (bulkResponse.hasFailures()) {
+			Map<String, String> failedDocuments = new HashMap<>();
+			for (BulkItemResponse item : bulkResponse.getItems()) {
+				if (item.isFailed())
+					failedDocuments.put(item.getId(), item.getFailureMessage());
+			}
+			throw new ElasticsearchException(
+					"Bulk operation has failures. Use ElasticsearchException.getFailedDocuments() for detailed messages ["
+							+ failedDocuments + "]",
+					failedDocuments);
+		}
+	}
+
+	/**
+	 * @param query
+	 * @param clazz
+	 * @deprecated index names and types should not be set in query
+	 */
+	@Deprecated
+	protected void setPersistentEntityIndexAndType(Query query, Class clazz) {
+		if (query.getIndices().isEmpty()) {
+			String[] indices = retrieveIndexNameFromPersistentEntity(clazz);
+
+			if (indices != null) {
+				query.addIndices(indices);
+			}
+		}
+		if (query.getTypes().isEmpty()) {
+			String[] types = retrieveTypeFromPersistentEntity(clazz);
+
+			if (types != null) {
+				query.addTypes(types);
+			}
+		}
+	}
+
+	private String[] retrieveIndexNameFromPersistentEntity(Class clazz) {
+		if (clazz != null) {
+			return new String[] { getPersistentEntityFor(clazz).getIndexName() };
+		}
+		return null;
+	}
+
+	private String[] retrieveTypeFromPersistentEntity(Class clazz) {
+		if (clazz != null) {
+			return new String[] { getPersistentEntityFor(clazz).getIndexType() };
+		}
+		return null;
+	}
+
+	@Override
+	public <T> List<Page<T>> queryForPage(List<? extends Query> queries, Class<T> clazz, IndexCoordinates index) {
+		MultiSearchRequest request = new MultiSearchRequest();
+		for (Query query : queries) {
+			request.add(requestFactory.searchRequest(query, clazz, index));
+		}
+		return doMultiSearch(queries, clazz, request);
+	}
+
+	@Override
+	public List<Page<?>> queryForPage(List<? extends Query> queries, List<Class<?>> classes, IndexCoordinates index) {
+		MultiSearchRequest request = new MultiSearchRequest();
+		Iterator<Class<?>> it = classes.iterator();
+		for (Query query : queries) {
+			request.add(requestFactory.searchRequest(query, it.next(), index));
+		}
+		return doMultiSearch(queries, classes, request);
+	}
+
+	private <T> List<Page<T>> doMultiSearch(List<? extends Query> queries, Class<T> clazz, MultiSearchRequest request) {
+		MultiSearchResponse.Item[] items = getMultiSearchResult(request);
+		List<Page<T>> res = new ArrayList<>(queries.size());
+		int c = 0;
+		for (Query query : queries) {
+			res.add(elasticsearchConverter.mapResults(SearchDocumentResponse.from(items[c++].getResponse()), clazz,
+					query.getPageable()));
+		}
+		return res;
+	}
+
+	private List<Page<?>> doMultiSearch(List<? extends Query> queries, List<Class<?>> classes,
+			MultiSearchRequest request) {
+		MultiSearchResponse.Item[] items = getMultiSearchResult(request);
+		List<Page<?>> res = new ArrayList<>(queries.size());
+		int c = 0;
+		Iterator<Class<?>> it = classes.iterator();
+		for (Query query : queries) {
+			res.add(elasticsearchConverter.mapResults(SearchDocumentResponse.from(items[c++].getResponse()), it.next(),
+					query.getPageable()));
+		}
+		return res;
+	}
+
+	@Override
+	public <T> boolean putMapping(Class<T> clazz) {
+		return putMapping(clazz, buildMapping(clazz));
+	}
+
+	@Override
+	public <T> boolean putMapping(Class<T> clazz, Object mapping) {
+		return putMapping(getIndexCoordinatesFor(clazz), mapping);
+	}
+
+	@Override
+	public <T> boolean putMapping(IndexCoordinates index, Class<T> clazz) {
+		return putMapping(index, buildMapping(clazz));
+	}
+
+	abstract protected MultiSearchResponse.Item[] getMultiSearchResult(MultiSearchRequest request);
+
+	protected void setPersistentEntityId(Object entity, String id) {
+
+		ElasticsearchPersistentEntity<?> persistentEntity = getPersistentEntityFor(entity.getClass());
+		ElasticsearchPersistentProperty idProperty = persistentEntity.getIdProperty();
+
+		// Only deal with text because ES generated Ids are strings !
+
+		if (idProperty != null && idProperty.getType().isAssignableFrom(String.class)) {
+			persistentEntity.getPropertyAccessor(entity).setProperty(idProperty, id);
+		}
+	}
 }