Add AfterLoad callback.

Original Pull Request #2039
Closes #2009

Author: sothawo

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

@@ -69,6 +69,7 @@
 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 @@ protected <T> Mono<T> maybeCallAfterConvert(T entity, Document document, IndexCo
 
 		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 Mono<T> toEntity(@Nullable Document document) {
 				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.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 @@ protected <T> T maybeCallbackAfterConvert(T entity, Document document, IndexCoor
 		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 T doWith(@Nullable Document document) {
 			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.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 @@ public SampleEntity onBeforeConvert(SampleEntity entity, IndexCoordinates index)
 				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 @@ void shouldApplyConversionResultToIndexQueryInBulkIndex() {
 		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 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.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 Mono<SampleEntity> onBeforeConvert(SampleEntity entity, IndexCoordinates
 			}
 		}
 
+		@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 @@ void shouldCallCallbackOnSave() {
 				.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;