Add AfterLoad callback.

Original Pull Request #2039
Closes #2009

src/main/java/org/springframework/data/elasticsearch/backend/elasticsearch7/ReactiveElasticsearchTemplate.java

@@ -69,6 +69,7 @@ import org.springframework.data.elasticsearch.core.convert.MappingElasticsearchC
 import org.springframework.data.elasticsearch.core.document.Document;
 import org.springframework.data.elasticsearch.core.document.SearchDocument;
 import org.springframework.data.elasticsearch.core.event.ReactiveAfterConvertCallback;
+import org.springframework.data.elasticsearch.core.event.ReactiveAfterLoadCallback;
 import org.springframework.data.elasticsearch.core.event.ReactiveAfterSaveCallback;
 import org.springframework.data.elasticsearch.core.event.ReactiveBeforeConvertCallback;
 import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
@@ -1159,6 +1160,17 @@ public class ReactiveElasticsearchTemplate implements ReactiveElasticsearchOpera
 
 		return Mono.just(entity);
 	}
+
+	protected <T> Mono<Document> maybeCallbackAfterLoad(Document document, Class<T> type,
+			IndexCoordinates indexCoordinates) {
+
+		if (entityCallbacks != null) {
+			return entityCallbacks.callback(ReactiveAfterLoadCallback.class, document, type, indexCoordinates);
+		}
+
+		return Mono.just(document);
+	}
+
 	// endregion
 
 	// region routing
@@ -1206,12 +1218,20 @@ public class ReactiveElasticsearchTemplate implements ReactiveElasticsearchOpera
 				return Mono.empty();
 			}
 
-			T entity = reader.read(type, document);
-			IndexedObjectInformation indexedObjectInformation = IndexedObjectInformation.of(
-					document.hasId() ? document.getId() : null, document.getSeqNo(), document.getPrimaryTerm(),
-					document.getVersion());
-			entity = updateIndexedObject(entity, indexedObjectInformation);
-			return maybeCallAfterConvert(entity, document, index);
+			return maybeCallbackAfterLoad(document, type, index) //
+					.flatMap(documentAfterLoad -> {
+
+						T entity = reader.read(type, documentAfterLoad);
+
+						IndexedObjectInformation indexedObjectInformation = IndexedObjectInformation.of( //
+								documentAfterLoad.hasId() ? documentAfterLoad.getId() : null, //
+								documentAfterLoad.getSeqNo(), //
+								documentAfterLoad.getPrimaryTerm(), //
+								documentAfterLoad.getVersion()); //
+						entity = updateIndexedObject(entity, indexedObjectInformation);
+
+						return maybeCallAfterConvert(entity, documentAfterLoad, index);
+					});
 		}
 	}
 

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

@@ -33,6 +33,7 @@ import org.springframework.data.elasticsearch.core.convert.ElasticsearchConverte
 import org.springframework.data.elasticsearch.core.convert.MappingElasticsearchConverter;
 import org.springframework.data.elasticsearch.core.document.Document;
 import org.springframework.data.elasticsearch.core.event.AfterConvertCallback;
+import org.springframework.data.elasticsearch.core.event.AfterLoadCallback;
 import org.springframework.data.elasticsearch.core.event.AfterSaveCallback;
 import org.springframework.data.elasticsearch.core.event.BeforeConvertCallback;
 import org.springframework.data.elasticsearch.core.mapping.ElasticsearchPersistentEntity;
@@ -690,6 +691,15 @@ public abstract class AbstractElasticsearchTemplate implements ElasticsearchOper
 		return entity;
 	}
 
+	protected <T> Document maybeCallbackAfterLoad(Document document, Class<T> type, IndexCoordinates indexCoordinates) {
+
+		if (entityCallbacks != null) {
+			return entityCallbacks.callback(AfterLoadCallback.class, document, type, indexCoordinates);
+		}
+
+		return document;
+	}
+
 	// endregion
 
 	protected void updateIndexedObjectsWithQueries(List<?> queries,
@@ -736,13 +746,18 @@ public abstract class AbstractElasticsearchTemplate implements ElasticsearchOper
 			if (document == null) {
 				return null;
 			}
+			Document documentAfterLoad = maybeCallbackAfterLoad(document, type, index);
 
-			T entity = reader.read(type, document);
-			IndexedObjectInformation indexedObjectInformation = IndexedObjectInformation.of(
-					document.hasId() ? document.getId() : null, document.getSeqNo(), document.getPrimaryTerm(),
-					document.getVersion());
+			T entity = reader.read(type, documentAfterLoad);
+
+			IndexedObjectInformation indexedObjectInformation = IndexedObjectInformation.of( //
+					documentAfterLoad.hasId() ? documentAfterLoad.getId() : null, //
+					documentAfterLoad.getSeqNo(), //
+					documentAfterLoad.getPrimaryTerm(), //
+					documentAfterLoad.getVersion()); //
 			entity = updateIndexedObject(entity, indexedObjectInformation);
-			return maybeCallbackAfterConvert(entity, document, index);
+
+			return maybeCallbackAfterConvert(entity, documentAfterLoad, index);
 		}
 	}
 

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

@@ -0,0 +1,42 @@
+/*
+ * Copyright 2021 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.elasticsearch.core.document.Document;
+import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
+import org.springframework.data.mapping.callback.EntityCallback;
+
+/**
+ * Callback being invoked after a {@link Document} is read from Elasticsearch and before it is converted into a domain
+ * object.
+ *
+ * @author Peter-Josef Meisch
+ * @since 4.4
+ * @see org.springframework.data.mapping.callback.EntityCallbacks
+ */
+@FunctionalInterface
+public interface AfterLoadCallback<T> extends EntityCallback<Document> {
+
+	/**
+	 * Entity callback method invoked after a domain object is materialized from a {@link Document}. Can return either the
+	 * same or a modified instance of the {@link Document} object.
+	 *
+	 * @param document the document.
+	 * @param indexCoordinates of the index the document was read from.
+	 * @return a possible modified or new {@link Document}.
+	 */
+	Document onAfterLoad(Document document, Class<T> type, IndexCoordinates indexCoordinates);
+}

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

@@ -0,0 +1,43 @@
+/*
+ * Copyright 2021 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.elasticsearch.core.document.Document;
+import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
+import org.springframework.data.mapping.callback.EntityCallback;
+
+/**
+ * Callback being invoked after a {@link Document} is read from Elasticsearch and before it is converted into a domain
+ * object.
+ *
+ * @author Peter-Josef Meisch
+ * @since 4.4
+ * @see org.springframework.data.mapping.callback.EntityCallbacks
+ */
+@FunctionalInterface
+public interface ReactiveAfterLoadCallback<T> extends EntityCallback<Document> {
+
+	/**
+	 * Entity callback method invoked after a domain object is materialized from a {@link Document}. Can return either the
+	 * same or a modified instance of the {@link Document} object.
+	 *
+	 * @param document the document.
+	 * @param indexCoordinates of the index the document was read from.
+	 * @return a possible modified or new {@link Document}.
+	 */
+	Publisher<Document> onAfterLoad(Document document, Class<T> type, IndexCoordinates indexCoordinates);
+}

src/test/java/org/springframework/data/elasticsearch/core/event/ElasticsearchOperationsCallbackIntegrationTests.java

@@ -29,6 +29,7 @@ import org.mockito.ArgumentCaptor;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.data.annotation.Id;
+import org.springframework.data.annotation.ReadOnlyProperty;
 import org.springframework.data.elasticsearch.annotations.Document;
 import org.springframework.data.elasticsearch.annotations.JoinTypeRelation;
 import org.springframework.data.elasticsearch.annotations.JoinTypeRelations;
@@ -73,6 +74,19 @@ abstract class ElasticsearchOperationsCallbackIntegrationTests {
 				return entity;
 			}
 		}
+
+		@Component
+		static class SampleEntityAfterLoadCallback implements AfterLoadCallback<SampleEntity> {
+
+			@Override
+			public org.springframework.data.elasticsearch.core.document.Document onAfterLoad(
+					org.springframework.data.elasticsearch.core.document.Document document, Class<SampleEntity> type,
+					IndexCoordinates indexCoordinates) {
+
+				document.put("className", document.get("_class"));
+				return document;
+			}
+		}
 	}
 
 	@BeforeEach
@@ -214,12 +228,30 @@ abstract class ElasticsearchOperationsCallbackIntegrationTests {
 		assertThat(capturedIndexQuery.getPrimaryTerm()).isEqualTo(seqNoPrimaryTerm.getPrimaryTerm());
 	}
 
+	@Test // #2009
+	@DisplayName("should invoke after load callback")
+	void shouldInvokeAfterLoadCallback() {
+
+		SampleEntity entity = new SampleEntity("1", "test");
+		operations.save(entity);
+
+		SampleEntity loaded = operations.get(entity.getId(), SampleEntity.class);
+
+		assertThat(loaded).isNotNull();
+		assertThat(loaded.className).isEqualTo(SampleEntity.class.getName());
+	}
+
 	@Document(indexName = INDEX)
 	static class SampleEntity {
-		@Nullable @Id private String id;
+		@Nullable
+		@Id private String id;
 		@Nullable private String text;
 
-		@Nullable @JoinTypeRelations(relations = {
+		@ReadOnlyProperty
+		@Nullable private String className;
+
+		@Nullable
+		@JoinTypeRelations(relations = {
 				@JoinTypeRelation(parent = "question", children = { "answer" }) }) private JoinField<String> joinField;
 
 		@Nullable private SeqNoPrimaryTerm seqNoPrimaryTerm;

src/test/java/org/springframework/data/elasticsearch/core/event/ReactiveElasticsearchOperationsCallbackTest.java

@@ -22,11 +22,13 @@ import reactor.test.StepVerifier;
 
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.context.annotation.Import;
 import org.springframework.data.annotation.Id;
+import org.springframework.data.annotation.ReadOnlyProperty;
 import org.springframework.data.elasticsearch.annotations.Document;
 import org.springframework.data.elasticsearch.core.ElasticsearchOperations;
 import org.springframework.data.elasticsearch.core.IndexOperations;
@@ -35,6 +37,7 @@ import org.springframework.data.elasticsearch.core.mapping.IndexCoordinates;
 import org.springframework.data.elasticsearch.junit.jupiter.ElasticsearchRestTemplateConfiguration;
 import org.springframework.data.elasticsearch.junit.jupiter.ReactiveElasticsearchRestTemplateConfiguration;
 import org.springframework.data.elasticsearch.junit.jupiter.SpringIntegrationTest;
+import org.springframework.lang.Nullable;
 import org.springframework.stereotype.Component;
 import org.springframework.test.context.ContextConfiguration;
 
@@ -49,6 +52,7 @@ public class ReactiveElasticsearchOperationsCallbackTest {
 	@Configuration
 	@Import({ ReactiveElasticsearchRestTemplateConfiguration.class, ElasticsearchRestTemplateConfiguration.class })
 	static class Config {
+
 		@Component
 		static class SampleEntityBeforeConvertCallback implements ReactiveBeforeConvertCallback<SampleEntity> {
 			@Override
@@ -58,6 +62,20 @@ public class ReactiveElasticsearchOperationsCallbackTest {
 			}
 		}
 
+		@Component
+		static class SampleEntityAfterLoadCallback
+				implements ReactiveAfterLoadCallback<ElasticsearchOperationsCallbackIntegrationTests.SampleEntity> {
+
+			@Override
+			public Mono<org.springframework.data.elasticsearch.core.document.Document> onAfterLoad(
+					org.springframework.data.elasticsearch.core.document.Document document,
+					Class<ElasticsearchOperationsCallbackIntegrationTests.SampleEntity> type, IndexCoordinates indexCoordinates) {
+
+				document.put("className", document.get("_class"));
+				return Mono.just(document);
+			}
+		}
+
 	}
 
 	@Autowired private ReactiveElasticsearchOperations operations;
@@ -88,11 +106,29 @@ public class ReactiveElasticsearchOperationsCallbackTest {
 				.verifyComplete();
 	}
 
+	@Test // #2009
+	@DisplayName("should invoke after load callback")
+	void shouldInvokeAfterLoadCallback() {
+
+		SampleEntity entity = new SampleEntity("1", "test");
+
+		operations.save(entity) //
+				.then(operations.get(entity.getId(), SampleEntity.class)) //
+				.as(StepVerifier::create) //
+				.consumeNextWith(loaded -> { //
+					assertThat(loaded).isNotNull(); //
+					assertThat(loaded.className).isEqualTo(SampleEntity.class.getName()); //
+				}).verifyComplete(); //
+	}
+
 	@Document(indexName = "test-operations-reactive-callback")
 	static class SampleEntity {
 		@Id private String id;
 		private String text;
 
+		@ReadOnlyProperty
+		@Nullable private String className;
+
 		public SampleEntity(String id, String text) {
 			this.id = id;
 			this.text = text;

src/test/java/org/springframework/data/elasticsearch/core/suggest/SuggestReactiveTemplateIntegrationTests.java

@@ -17,6 +17,7 @@ package org.springframework.data.elasticsearch.core.suggest;
 
 import static org.assertj.core.api.Assertions.*;
 
+import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
 import java.util.ArrayList;
@@ -54,6 +55,7 @@ import org.springframework.lang.Nullable;
 @SuppressWarnings("SpringJavaAutowiredMembersInspection")
 @SpringIntegrationTest
 public class SuggestReactiveTemplateIntegrationTests {
+
 	@Configuration
 	@Import({ ReactiveElasticsearchRestTemplateConfiguration.class })
 	static class Config {
@@ -86,32 +88,34 @@ public class SuggestReactiveTemplateIntegrationTests {
 	@DisplayName("should find suggestions for given prefix completion")
 	void shouldFindSuggestionsForGivenPrefixCompletion() {
 
-		loadCompletionObjectEntities();
+		loadCompletionObjectEntities().map(unused -> {
 
-		NativeSearchQuery query = new NativeSearchQueryBuilder().withSuggestBuilder(new SuggestBuilder()
-				.addSuggestion("test-suggest", SuggestBuilders.completionSuggestion("suggest").prefix("m", Fuzziness.AUTO)))
-				.build();
+			NativeSearchQuery query = new NativeSearchQueryBuilder().withSuggestBuilder(new SuggestBuilder()
+					.addSuggestion("test-suggest", SuggestBuilders.completionSuggestion("suggest").prefix("m", Fuzziness.AUTO)))
+					.build();
 
-		operations.suggest(query, CompletionEntity.class) //
-				.as(StepVerifier::create) //
-				.assertNext(suggest -> {
-					Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> suggestion = suggest
-							.getSuggestion("test-suggest");
-					assertThat(suggestion).isNotNull();
-					assertThat(suggestion).isInstanceOf(CompletionSuggestion.class);
-					// noinspection unchecked
-					List<CompletionSuggestion.Entry.Option<CompletionIntegrationTests.AnnotatedCompletionEntity>> options = ((CompletionSuggestion<CompletionIntegrationTests.AnnotatedCompletionEntity>) suggestion)
-							.getEntries().get(0).getOptions();
-					assertThat(options).hasSize(2);
-					assertThat(options.get(0).getText()).isIn("Marchand", "Mohsin");
-					assertThat(options.get(1).getText()).isIn("Marchand", "Mohsin");
+			operations.suggest(query, CompletionEntity.class) //
+					.as(StepVerifier::create) //
+					.assertNext(suggest -> {
+						Suggest.Suggestion<? extends Suggest.Suggestion.Entry<? extends Suggest.Suggestion.Entry.Option>> suggestion = suggest
+								.getSuggestion("test-suggest");
+						assertThat(suggestion).isNotNull();
+						assertThat(suggestion).isInstanceOf(CompletionSuggestion.class);
+						// noinspection unchecked
+						List<CompletionSuggestion.Entry.Option<CompletionIntegrationTests.AnnotatedCompletionEntity>> options = ((CompletionSuggestion<CompletionIntegrationTests.AnnotatedCompletionEntity>) suggestion)
+								.getEntries().get(0).getOptions();
+						assertThat(options).hasSize(2);
+						assertThat(options.get(0).getText()).isIn("Marchand", "Mohsin");
+						assertThat(options.get(1).getText()).isIn("Marchand", "Mohsin");
 
-				}) //
-				.verifyComplete();
+					}) //
+					.verifyComplete();
+			return Mono.empty();
+		});
 	}
 
 	// region helper functions
-	private void loadCompletionObjectEntities() {
+	private Mono<Void> loadCompletionObjectEntities() {
 
 		CompletionEntity rizwan_idrees = new CompletionEntityBuilder("1").name("Rizwan Idrees")
 				.suggest(new String[] { "Rizwan Idrees" }).build();
@@ -124,7 +128,7 @@ public class SuggestReactiveTemplateIntegrationTests {
 		List<CompletionEntity> entities = new ArrayList<>(
 				Arrays.asList(rizwan_idrees, franck_marchand, mohsin_husen, artur_konczak));
 		IndexCoordinates index = IndexCoordinates.of(indexNameProvider.indexName());
-		operations.saveAll(entities, index).blockLast();
+		return operations.saveAll(entities, index).then();
 	}
 	// endregion
 
@@ -132,11 +136,13 @@ public class SuggestReactiveTemplateIntegrationTests {
 	@Document(indexName = "#{@indexNameProvider.indexName()}")
 	static class CompletionEntity {
 
-		@Nullable @Id private String id;
+		@Nullable
+		@Id private String id;
 
 		@Nullable private String name;
 
-		@Nullable @CompletionField(maxInputLength = 100) private Completion suggest;
+		@Nullable
+		@CompletionField(maxInputLength = 100) private Completion suggest;
 
 		private CompletionEntity() {}