DATAES-771 - Add after-save entity callbacks support.

This change enables to get callbacks when an entity gets saved to Elasticsearch.

Author: rpuch

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

@@ -22,6 +22,7 @@
 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.event.AfterSaveCallback;
 import org.springframework.data.elasticsearch.core.event.BeforeConvertCallback;
 import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
 import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
@@ -45,6 +46,7 @@
  *
  * @author Sascha Woo
  * @author Peter-Josef Meisch
+ * @author Roman Puchkovskiy
  */
 public abstract class AbstractElasticsearchTemplate implements ElasticsearchOperations, ApplicationContextAware {
 
@@ -117,8 +119,13 @@ public <T> T save(T entity, IndexCoordinates index) {
 		Assert.notNull(entity, "entity must not be null");
 		Assert.notNull(index, "index must not be null");
 
-		index(getIndexQuery(entity), index);
-		return entity;
+		IndexQuery query = getIndexQuery(entity);
+		index(query, index);
+
+		// suppressing because it's either entity itself or something of a correct type returned by an entity callback
+		@SuppressWarnings("unchecked")
+		T castResult = (T) query.getObject();
+		return castResult;
 	}
 
 	@Override
@@ -151,7 +158,10 @@ public <T> Iterable<T> save(Iterable<T> entities, IndexCoordinates index) {
 			});
 		}
 
-		return entities;
+		return indexQueries.stream()
+				.map(IndexQuery::getObject)
+				.map(entity -> (T) entity)
+				.collect(Collectors.toList());
 	}
 
 	@Override
@@ -455,11 +465,39 @@ protected void maybeCallbackBeforeConvertWithQuery(Object query) {
 	}
 
 	// this can be called with either a List<IndexQuery> or a List<UpdateQuery>; these query classes
-	// don't have a common bas class, therefore the List<?> argument
+	// don't have a common base class, therefore the List<?> argument
 	protected void maybeCallbackBeforeConvertWithQueries(List<?> queries) {
 		queries.forEach(this::maybeCallbackBeforeConvertWithQuery);
 	}
 
+	protected <T> T maybeCallbackAfterSave(T entity) {
+
+		if (entityCallbacks != null) {
+			return entityCallbacks.callback(AfterSaveCallback.class, entity);
+		}
+
+		return entity;
+	}
+
+	protected void maybeCallbackAfterSaveWithQuery(Object query) {
+
+		if (query instanceof IndexQuery) {
+			IndexQuery indexQuery = (IndexQuery) query;
+			Object queryObject = indexQuery.getObject();
+
+			if (queryObject != null) {
+				queryObject = maybeCallbackAfterSave(queryObject);
+				indexQuery.setObject(queryObject);
+			}
+		}
+	}
+
+	// this can be called with either a List<IndexQuery> or a List<UpdateQuery>; these query classes
+	// don't have a common base class, therefore the List<?> argument
+	protected void maybeCallbackAfterSaveWithQueries(List<?> queries) {
+		queries.forEach(this::maybeCallbackAfterSaveWithQuery);
+	}
+
 	// endregion
 
 }

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

@@ -145,6 +145,9 @@ public String index(IndexQuery query, IndexCoordinates index) {
 		if (queryObject != null) {
 			setPersistentEntityId(queryObject, documentId);
 		}
+
+		maybeCallbackAfterSaveWithQuery(query);
+
 		return documentId;
 	}
 
@@ -226,7 +229,10 @@ public UpdateResponse update(UpdateQuery query, IndexCoordinates index) {
 	private List<String> doBulkOperation(List<?> queries, BulkOptions bulkOptions, IndexCoordinates index) {
 		maybeCallbackBeforeConvertWithQueries(queries);
 		BulkRequest bulkRequest = requestFactory.bulkRequest(queries, bulkOptions, index);
-		return checkForBulkOperationFailure(execute(client -> client.bulk(bulkRequest, RequestOptions.DEFAULT)));
+		List<String> ids = checkForBulkOperationFailure(execute(
+				client -> client.bulk(bulkRequest, RequestOptions.DEFAULT)));
+		maybeCallbackAfterSaveWithQueries(queries);
+		return ids;
 	}
 	// endregion
 

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

@@ -31,14 +31,12 @@
 import org.elasticsearch.action.search.MultiSearchResponse;
 import org.elasticsearch.action.search.SearchRequestBuilder;
 import org.elasticsearch.action.search.SearchResponse;
-import org.elasticsearch.action.search.SearchScrollRequestBuilder;
 import org.elasticsearch.action.update.UpdateRequestBuilder;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.search.suggest.SuggestBuilder;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.springframework.data.domain.Pageable;
 import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverter;
 import org.springframework.data.elasticsearch.core.document.DocumentAdapters;
 import org.springframework.data.elasticsearch.core.document.SearchDocumentResponse;
@@ -79,6 +77,7 @@
  * @author Martin Choraine
  * @author Farid Azaza
  * @author Gyula Attila Csorogi
+ * @author Roman Puchkovskiy
  * @deprecated as of 4.0
  */
 @Deprecated
@@ -153,6 +152,8 @@ public String index(IndexQuery query, IndexCoordinates index) {
 			setPersistentEntityId(queryObject, documentId);
 		}
 
+		maybeCallbackAfterSaveWithQuery(query);
+
 		return documentId;
 	}
 
@@ -188,7 +189,11 @@ public List<String> bulkIndex(List<IndexQuery> queries, BulkOptions bulkOptions,
 		Assert.notNull(queries, "List of IndexQuery must not be null");
 		Assert.notNull(bulkOptions, "BulkOptions must not be null");
 
-		return doBulkOperation(queries, bulkOptions, index);
+		List<String> ids = doBulkOperation(queries, bulkOptions, index);
+
+		maybeCallbackAfterSaveWithQueries(queries);
+
+		return ids;
 	}
 
 	@Override

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

@@ -70,6 +70,7 @@
 import org.springframework.data.elasticsearch.core.document.Document;
 import org.springframework.data.elasticsearch.core.document.DocumentAdapters;
 import org.springframework.data.elasticsearch.core.document.SearchDocument;
+import org.springframework.data.elasticsearch.core.event.ReactiveAfterSaveCallback;
 import org.springframework.data.elasticsearch.core.event.ReactiveBeforeConvertCallback;
 import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
 import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentProperty;
@@ -97,6 +98,7 @@
  * @author Peter-Josef Meisch
  * @author Mathias Teier
  * @author Aleksei Arsenev
+ * @author Roman Puchkovskiy
  * @since 3.2
  */
 public class ReactiveElasticsearchTemplate implements ReactiveElasticsearchOperations, ApplicationContextAware {
@@ -185,7 +187,8 @@ public <T> Mono<T> save(T entity, IndexCoordinates index) {
 		return doIndex(entity, adaptableEntity, index) //
 				.map(it -> {
 					return adaptableEntity.populateIdIfNecessary(it.getId());
-				});
+				})
+				.flatMap(this::maybeCallAfterSave);
 	}
 
 	@Override
@@ -213,11 +216,11 @@ public <T> Flux<T> saveAll(Mono<? extends Collection<? extends T>> entities, Ind
 					.map(e -> getIndexQuery(e.getBean(), e)) //
 					.collect(Collectors.toList());
 			return doBulkOperation(indexRequests, BulkOptions.defaultOptions(), index) //
-					.map(bulkItemResponse -> {
+					.flatMap(bulkItemResponse -> {
 
 						AdaptibleEntity<? extends T> mappedEntity = iterator.next();
 						mappedEntity.populateIdIfNecessary(bulkItemResponse.getResponse().getId());
-						return mappedEntity.getBean();
+						return maybeCallAfterSave(mappedEntity.getBean());
 					});
 		});
 	}
@@ -882,5 +885,14 @@ protected <T> Mono<T> maybeCallBeforeConvert(T entity) {
 
 		return Mono.just(entity);
 	}
+
+	protected <T> Mono<T> maybeCallAfterSave(T entity) {
+
+		if (null != entityCallbacks) {
+			return entityCallbacks.callback(ReactiveAfterSaveCallback.class, entity);
+		}
+
+		return Mono.just(entity);
+	}
 	// endregion
 }

src/main/java/org/springframework/data/elasticsearch/core/event/AfterSaveCallback.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2020 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.core.event;
+
+import org.springframework.data.mapping.callback.EntityCallback;
+import org.springframework.data.mapping.callback.EntityCallbacks;
+
+/**
+ * Entity callback triggered after save of an entity.
+ *
+ * @author Roman Puchkovskiy
+ * @since 4.0
+ * @see EntityCallbacks
+ */
+@FunctionalInterface
+public interface AfterSaveCallback<T> extends EntityCallback<T> {
+
+	/**
+	 * Entity callback method invoked after a domain object is saved. Can return either the same or a modified
+	 * instance of the domain object.
+	 *
+	 * @param entity the domain object that was saved.
+	 * @return the domain object that was persisted.
+	 */
+	T onAfterSave(T entity);
+}

src/main/java/org/springframework/data/elasticsearch/core/event/ReactiveAfterSaveCallback.java

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2020 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.core.event;
+
+import org.reactivestreams.Publisher;
+import org.springframework.data.mapping.callback.EntityCallback;
+import org.springframework.data.mapping.callback.ReactiveEntityCallbacks;
+
+/**
+ * Entity callback triggered after save of an entity.
+ *
+ * @author Roman Puchkovskiy
+ * @since 4.0
+ * @see ReactiveEntityCallbacks
+ */
+@FunctionalInterface
+public interface ReactiveAfterSaveCallback<T> extends EntityCallback<T> {
+
+	/**
+	 * Entity callback method invoked after a domain object is saved. Can return either the same or a modified
+	 * instance of the domain object.
+	 *
+	 * @param entity the domain object that was saved.
+	 * @return a {@link Publisher} emitting the domain object to be returned to the caller.
+	 */
+	Publisher<T> onAfterSave(T entity);
+}