All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.springframework.data.elasticsearch.client.elc.ReactiveElasticsearchTemplate Maven / Gradle / Ivy

Go to download

Spring Data Implementation for Elasticsearch

There is a newer version: 5.3.3
Show newest version
/*
 * Copyright 2021-2023 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.springframework.data.elasticsearch.client.elc;

import static co.elastic.clients.util.ApiTypeHelper.DANGEROUS_disableRequiredPropertiesCheck;
import static org.springframework.data.elasticsearch.client.elc.TypeUtils.result;

import co.elastic.clients.elasticsearch._types.Result;
import co.elastic.clients.elasticsearch.core.*;
import co.elastic.clients.elasticsearch.core.bulk.BulkResponseItem;
import co.elastic.clients.elasticsearch.core.get.GetResult;
import co.elastic.clients.elasticsearch.core.search.ResponseBody;
import co.elastic.clients.json.JsonpMapper;
import co.elastic.clients.transport.Version;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.util.function.Tuple2;

import java.time.Duration;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.function.Function;
import java.util.stream.Collectors;

import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.domain.Sort;
import org.springframework.data.elasticsearch.BulkFailureException;
import org.springframework.data.elasticsearch.NoSuchIndexException;
import org.springframework.data.elasticsearch.UncategorizedElasticsearchException;
import org.springframework.data.elasticsearch.client.UnsupportedBackendOperation;
import org.springframework.data.elasticsearch.core.*;
import org.springframework.data.elasticsearch.core.cluster.ReactiveClusterOperations;
import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
import org.springframework.data.elasticsearch.core.document.Document;
import org.springframework.data.elasticsearch.core.document.SearchDocument;
import org.springframework.data.elasticsearch.core.document.SearchDocumentResponse;
import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
import org.springframework.data.elasticsearch.core.query.*;
import org.springframework.data.elasticsearch.core.query.UpdateResponse;
import org.springframework.data.elasticsearch.core.reindex.ReindexRequest;
import org.springframework.data.elasticsearch.core.reindex.ReindexResponse;
import org.springframework.data.elasticsearch.core.script.Script;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;

/**
 * Implementation of {@link org.springframework.data.elasticsearch.core.ReactiveElasticsearchOperations} using the new
 * Elasticsearch client.
 *
 * @author Peter-Josef Meisch
 * @since 4.4
 */
public class ReactiveElasticsearchTemplate extends AbstractReactiveElasticsearchTemplate {

	private static final Logger LOGGER = LoggerFactory.getLogger(ReactiveElasticsearchTemplate.class);

	private final ReactiveElasticsearchClient client;
	private final RequestConverter requestConverter;
	private final ResponseConverter responseConverter;
	private final JsonpMapper jsonpMapper;
	private final ElasticsearchExceptionTranslator exceptionTranslator;

	public ReactiveElasticsearchTemplate(ReactiveElasticsearchClient client, ElasticsearchConverter converter) {
		super(converter);

		Assert.notNull(client, "client must not be null");

		this.client = client;
		this.jsonpMapper = client._transport().jsonpMapper();
		requestConverter = new RequestConverter(converter, jsonpMapper);
		responseConverter = new ResponseConverter(jsonpMapper);
		exceptionTranslator = new ElasticsearchExceptionTranslator(jsonpMapper);
	}

	// region Document operations
	@Override
	protected  Mono> doIndex(T entity, IndexCoordinates index) {

		IndexRequest indexRequest = requestConverter.documentIndexRequest(getIndexQuery(entity), index,
				getRefreshPolicy());
		return Mono.just(entity) //
				.zipWith(//
						Mono.from(execute((ClientCallback>) client -> client.index(indexRequest))) //
								.map(indexResponse -> new IndexResponseMetaData(indexResponse.id(), //
										indexResponse.index(), //
										indexResponse.seqNo(), //
										indexResponse.primaryTerm(), //
										indexResponse.version() //
								)));
	}

	@Override
	public  Flux saveAll(Mono> entitiesPublisher, IndexCoordinates index) {

		Assert.notNull(entitiesPublisher, "entitiesPublisher must not be null!");

		return entitiesPublisher //
				.flatMapMany(entities -> Flux.fromIterable(entities) //
						.concatMap(entity -> maybeCallbackBeforeConvert(entity, index)) //
				).collectList() //
				.map(Entities::new) //
				.flatMapMany(entities -> {

					if (entities.isEmpty()) {
						return Flux.empty();
					}

					return doBulkOperation(entities.indexQueries(), BulkOptions.defaultOptions(), index)//
							.index() //
							.flatMap(indexAndResponse -> {
								T savedEntity = entities.entityAt(indexAndResponse.getT1());
								BulkResponseItem response = indexAndResponse.getT2();
								updateIndexedObject(savedEntity, new IndexedObjectInformation( //
										response.id(), //
										response.index(), //
										response.seqNo(), //
										response.primaryTerm(), //
										response.version()));
								return maybeCallbackAfterSave(savedEntity, index);
							});
				});
	}

	@Override
	protected Mono doExists(String id, IndexCoordinates index) {

		Assert.notNull(id, "id must not be null");
		Assert.notNull(index, "index must not be null");

		GetRequest getRequest = requestConverter.documentGetRequest(id, routingResolver.getRouting(), index, true);

		return Mono.from(execute(
				((ClientCallback>>) client -> client.get(getRequest, EntityAsMap.class))))
				.map(GetResult::found) //
				.onErrorReturn(NoSuchIndexException.class, false);
	}

	@Override
	public Mono delete(Query query, Class entityType, IndexCoordinates index) {

		Assert.notNull(query, "query must not be null");

		DeleteByQueryRequest request = requestConverter.documentDeleteByQueryRequest(query, routingResolver.getRouting(),
				entityType, index, getRefreshPolicy());
		return Mono
				.from(execute((ClientCallback>) client -> client.deleteByQuery(request)))
				.map(responseConverter::byQueryResponse);
	}

	@Override
	public  Mono get(String id, Class entityType, IndexCoordinates index) {

		Assert.notNull(id, "id must not be null");
		Assert.notNull(entityType, "entityType must not be null");
		Assert.notNull(index, "index must not be null");

		GetRequest getRequest = requestConverter.documentGetRequest(id, routingResolver.getRouting(), index, false);

		Mono> getResponse = Mono.from(execute(
				(ClientCallback>>) client -> client.get(getRequest, EntityAsMap.class)));

		ReadDocumentCallback callback = new ReadDocumentCallback<>(converter, entityType, index);
		return getResponse.flatMap(response -> callback.toEntity(DocumentAdapters.from(response)));
	}

	@Override
	public Mono reindex(ReindexRequest reindexRequest) {

		Assert.notNull(reindexRequest, "reindexRequest must not be null");

		co.elastic.clients.elasticsearch.core.ReindexRequest reindexRequestES = requestConverter.reindex(reindexRequest,
				true);

		return Mono.from(execute( //
				(ClientCallback>) client -> client
						.reindex(reindexRequestES)))
				.map(responseConverter::reindexResponse);
	}

	@Override
	public Mono submitReindex(ReindexRequest reindexRequest) {

		Assert.notNull(reindexRequest, "reindexRequest must not be null");

		co.elastic.clients.elasticsearch.core.ReindexRequest reindexRequestES = requestConverter.reindex(reindexRequest,
				false);

		return Mono.from(execute( //
				(ClientCallback>) client -> client
						.reindex(reindexRequestES)))
				.flatMap(response -> (response.task() == null)
						? Mono.error(
								new UnsupportedBackendOperation("ElasticsearchClient did not return a task id on submit request"))
						: Mono.just(response.task()));
	}

	@Override
	public Mono update(UpdateQuery updateQuery, IndexCoordinates index) {

		Assert.notNull(updateQuery, "UpdateQuery must not be null");
		Assert.notNull(index, "Index must not be null");

		UpdateRequest request = requestConverter.documentUpdateRequest(updateQuery, index, getRefreshPolicy(),
				routingResolver.getRouting());

		return Mono.from(execute(
				(ClientCallback>>) client -> client
						.update(request, Document.class)))
				.flatMap(response -> {
					UpdateResponse.Result result = result(response.result());
					return result == null ? Mono.empty() : Mono.just(UpdateResponse.of(result));
				});
	}

	@Override
	public Mono updateByQuery(UpdateQuery updateQuery, IndexCoordinates index) {
		throw new UnsupportedOperationException("not implemented");
	}

	@Override
	public Mono bulkUpdate(List queries, BulkOptions bulkOptions, IndexCoordinates index) {

		Assert.notNull(queries, "List of UpdateQuery must not be null");
		Assert.notNull(bulkOptions, "BulkOptions must not be null");
		Assert.notNull(index, "Index must not be null");

		return doBulkOperation(queries, bulkOptions, index).then();
	}

	private Flux doBulkOperation(List queries, BulkOptions bulkOptions, IndexCoordinates index) {

		BulkRequest bulkRequest = requestConverter.documentBulkRequest(queries, bulkOptions, index, getRefreshPolicy());
		return client.bulk(bulkRequest)
				.onErrorMap(e -> new UncategorizedElasticsearchException("Error executing bulk request", e))
				.flatMap(this::checkForBulkOperationFailure) //
				.flatMapMany(response -> Flux.fromIterable(response.items()));

	}

	private Mono checkForBulkOperationFailure(BulkResponse bulkResponse) {

		if (bulkResponse.errors()) {
			Map failedDocuments = new HashMap<>();

			for (BulkResponseItem item : bulkResponse.items()) {

				if (item.error() != null) {
					failedDocuments.put(item.id(), item.error().reason());
				}
			}
			BulkFailureException exception = new BulkFailureException(
					"Bulk operation has failures. Use ElasticsearchException.getFailedDocuments() for detailed messages ["
							+ failedDocuments + ']',
					failedDocuments);
			return Mono.error(exception);
		} else {
			return Mono.just(bulkResponse);
		}
	}

	@Override
	protected Mono doDeleteById(String id, @Nullable String routing, IndexCoordinates index) {

		Assert.notNull(id, "id must not be null");
		Assert.notNull(index, "index must not be null");

		return Mono.defer(() -> {
			DeleteRequest deleteRequest = requestConverter.documentDeleteRequest(id, routing, index, getRefreshPolicy());
			return doDelete(deleteRequest);
		});
	}

	private Mono doDelete(DeleteRequest request) {

		return Mono.from(execute((ClientCallback>) client -> client.delete(request))) //
				.flatMap(deleteResponse -> {
					if (deleteResponse.result() == Result.NotFound) {
						return Mono.empty();
					}
					return Mono.just(deleteResponse.id());
				}).onErrorResume(NoSuchIndexException.class, it -> Mono.empty());
	}

	@Override
	public  Flux> multiGet(Query query, Class clazz, IndexCoordinates index) {

		Assert.notNull(query, "query must not be null");
		Assert.notNull(clazz, "clazz must not be null");

		MgetRequest request = requestConverter.documentMgetRequest(query, clazz, index);

		ReadDocumentCallback callback = new ReadDocumentCallback<>(converter, clazz, index);

		Publisher> response = execute(
				(ClientCallback>>) client -> client.mget(request, EntityAsMap.class));

		return Mono.from(response)//
				.flatMapMany(it -> Flux.fromIterable(DocumentAdapters.from(it))) //
				.flatMap(multiGetItem -> {
					if (multiGetItem.isFailed()) {
						return Mono.just(MultiGetItem.of(null, multiGetItem.getFailure()));
					} else {
						return callback.toEntity(multiGetItem.getItem()) //
								.map(t -> MultiGetItem.of(t, multiGetItem.getFailure()));
					}
				});
	}

	// endregion

	@Override
	protected ReactiveElasticsearchTemplate doCopy() {
		return new ReactiveElasticsearchTemplate(client, converter);
	}

	// region search operations

	@Override
	protected Flux doFind(Query query, Class clazz, IndexCoordinates index) {

		Assert.notNull(query, "query must not be null");
		Assert.notNull(clazz, "clazz must not be null");
		Assert.notNull(index, "index must not be null");

		if (query instanceof SearchTemplateQuery searchTemplateQuery) {
			return Flux.defer(() -> doSearch(searchTemplateQuery, clazz, index));
		} else {
			return Flux.defer(() -> {
				boolean queryIsUnbounded = !(query.getPageable().isPaged() || query.isLimiting());
				return queryIsUnbounded ? doFindUnbounded(query, clazz, index) : doFindBounded(query, clazz, index);
			});
		}
	}

	private Flux doFindUnbounded(Query query, Class clazz, IndexCoordinates index) {

		if (query instanceof BaseQuery baseQuery) {
			var pitKeepAlive = Duration.ofMinutes(5);
			// setup functions for Flux.usingWhen()
			Mono resourceSupplier = openPointInTime(index, pitKeepAlive, true)
					.map(pit -> new PitSearchAfter(baseQuery, pit));

			Function> asyncComplete = this::cleanupPit;

			BiFunction> asyncError = (psa, ex) -> {
				if (LOGGER.isErrorEnabled()) {
					LOGGER.error(String.format("Error during pit/search_after"), ex);
				}
				return cleanupPit(psa);
			};

			Function> asyncCancel = psa -> {
				if (LOGGER.isWarnEnabled()) {
					LOGGER.warn(String.format("pit/search_after was cancelled"));
				}
				return cleanupPit(psa);
			};

			Function>> resourceClosure = psa -> {

				baseQuery.setPointInTime(new Query.PointInTime(psa.getPit(), pitKeepAlive));
				baseQuery.addSort(Sort.by("_shard_doc"));
				SearchRequest firstSearchRequest = requestConverter.searchRequest(baseQuery, routingResolver.getRouting(),
						clazz, index, false, true);

				return Mono.from(execute((ClientCallback>>) client -> client
						.search(firstSearchRequest, EntityAsMap.class))).expand(entityAsMapSearchResponse -> {

							var hits = entityAsMapSearchResponse.hits().hits();
							if (CollectionUtils.isEmpty(hits)) {
								return Mono.empty();
							}

							List sortOptions = hits.get(hits.size() - 1).sort().stream().map(TypeUtils::toObject)
									.collect(Collectors.toList());
							baseQuery.setSearchAfter(sortOptions);
							SearchRequest followSearchRequest = requestConverter.searchRequest(baseQuery,
									routingResolver.getRouting(), clazz, index, false, true);
							return Mono.from(execute((ClientCallback>>) client -> client
									.search(followSearchRequest, EntityAsMap.class)));
						});

			};

			Flux> searchResponses = Flux.usingWhen(resourceSupplier, resourceClosure, asyncComplete,
					asyncError, asyncCancel);
			return searchResponses.flatMapIterable(entityAsMapSearchResponse -> entityAsMapSearchResponse.hits().hits())
					.map(entityAsMapHit -> DocumentAdapters.from(entityAsMapHit, jsonpMapper));
		} else {
			return Flux.error(new IllegalArgumentException("Query must be derived from BaseQuery"));
		}
	}

	private Publisher cleanupPit(PitSearchAfter psa) {
		var baseQuery = psa.getBaseQuery();
		baseQuery.setPointInTime(null);
		baseQuery.setSearchAfter(null);
		baseQuery.setSort(psa.getSort());
		var pit = psa.getPit();
		return StringUtils.hasText(pit) ? closePointInTime(pit) : Mono.empty();
	}

	static private class PitSearchAfter {
		private final BaseQuery baseQuery;
		@Nullable private final Sort sort;
		private final String pit;

		PitSearchAfter(BaseQuery baseQuery, String pit) {
			this.baseQuery = baseQuery;
			this.sort = baseQuery.getSort();
			this.pit = pit;
		}

		public BaseQuery getBaseQuery() {
			return baseQuery;
		}

		@Nullable
		public Sort getSort() {
			return sort;
		}

		public String getPit() {
			return pit;
		}
	}

	@Override
	protected Mono doCount(Query query, Class entityType, IndexCoordinates index) {

		Assert.notNull(query, "query must not be null");
		Assert.notNull(index, "index must not be null");

		SearchRequest searchRequest = requestConverter.searchRequest(query, routingResolver.getRouting(), entityType, index,
				true);

		return Mono
				.from(execute((ClientCallback>>) client -> client.search(searchRequest,
						EntityAsMap.class)))
				.map(searchResponse -> searchResponse.hits().total() != null ? searchResponse.hits().total().value() : 0L);
	}

	private Flux doFindBounded(Query query, Class clazz, IndexCoordinates index) {

		SearchRequest searchRequest = requestConverter.searchRequest(query, routingResolver.getRouting(), clazz, index,
				false, false);

		return Mono
				.from(execute((ClientCallback>>) client -> client.search(searchRequest,
						EntityAsMap.class))) //
				.flatMapIterable(entityAsMapSearchResponse -> entityAsMapSearchResponse.hits().hits()) //
				.map(entityAsMapHit -> DocumentAdapters.from(entityAsMapHit, jsonpMapper));
	}

	private Flux doSearch(SearchTemplateQuery query, Class clazz, IndexCoordinates index) {

		var request = requestConverter.searchTemplate(query, routingResolver.getRouting(), index);

		return Mono
				.from(execute((ClientCallback>>) client -> client
						.searchTemplate(request, EntityAsMap.class))) //
				.flatMapIterable(entityAsMapSearchResponse -> entityAsMapSearchResponse.hits().hits()) //
				.map(entityAsMapHit -> DocumentAdapters.from(entityAsMapHit, jsonpMapper));
	}

	@Override
	protected  Mono doFindForResponse(Query query, Class clazz, IndexCoordinates index) {

		Assert.notNull(query, "query must not be null");
		Assert.notNull(index, "index must not be null");

		SearchRequest searchRequest = requestConverter.searchRequest(query, routingResolver.getRouting(), clazz, index,
				false);

		// noinspection unchecked
		SearchDocumentCallback callback = new ReadSearchDocumentCallback<>((Class) clazz, index);
		SearchDocumentResponse.EntityCreator entityCreator = searchDocument -> callback.toEntity(searchDocument)
				.toFuture();

		return Mono
				.from(execute((ClientCallback>>) client -> client.search(searchRequest,
						EntityAsMap.class)))
				.map(searchResponse -> SearchDocumentResponseBuilder.from(searchResponse, entityCreator, jsonpMapper));
	}

	@Override
	public Flux> aggregate(Query query, Class entityType, IndexCoordinates index) {

		return doFindForResponse(query, entityType, index).flatMapMany(searchDocumentResponse -> {
			ElasticsearchAggregations aggregations = (ElasticsearchAggregations) searchDocumentResponse.getAggregations();
			return aggregations == null ? Flux.empty() : Flux.fromIterable(aggregations.aggregations());
		});
	}

	@Override
	public Mono openPointInTime(IndexCoordinates index, Duration keepAlive, Boolean ignoreUnavailable) {

		Assert.notNull(index, "index must not be null");
		Assert.notNull(keepAlive, "keepAlive must not be null");
		Assert.notNull(ignoreUnavailable, "ignoreUnavailable must not be null");

		var request = requestConverter.searchOpenPointInTimeRequest(index, keepAlive, ignoreUnavailable);
		return Mono
				.from(execute((ClientCallback>) client -> client.openPointInTime(request)))
				.map(OpenPointInTimeResponse::id);
	}

	@Override
	public Mono closePointInTime(String pit) {

		Assert.notNull(pit, "pit must not be null");

		ClosePointInTimeRequest request = requestConverter.searchClosePointInTime(pit);
		return Mono
				.from(execute((ClientCallback>) client -> client.closePointInTime(request)))
				.map(ClosePointInTimeResponse::succeeded);
	}

	// endregion

	// region script operations
	@Override
	public Mono putScript(Script script) {

		Assert.notNull(script, "script must not be null");

		var request = requestConverter.scriptPut(script);
		return Mono.from(execute((ClientCallback>) client -> client.putScript(request)))
				.map(PutScriptResponse::acknowledged);
	}

	@Override
	public Mono