Callbacks refactoring

src/main/java/org/springframework/data/couchbase/core/AbstractTemplateSupport.java

@@ -15,11 +15,8 @@
  */
 package org.springframework.data.couchbase.core;
 
-import java.lang.reflect.InaccessibleObjectException;
-import java.util.Map;
-import java.util.Set;
-
 import com.couchbase.client.core.annotation.Stability;
+import com.couchbase.client.core.error.CouchbaseException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.context.ApplicationContext;
@@ -29,7 +26,6 @@
 import org.springframework.data.couchbase.core.mapping.CouchbaseDocument;
 import org.springframework.data.couchbase.core.mapping.CouchbasePersistentEntity;
 import org.springframework.data.couchbase.core.mapping.CouchbasePersistentProperty;
-import org.springframework.data.couchbase.core.mapping.event.AfterSaveEvent;
 import org.springframework.data.couchbase.core.mapping.event.CouchbaseMappingEvent;
 import org.springframework.data.couchbase.core.support.TemplateUtils;
 import org.springframework.data.couchbase.repository.support.MappingCouchbaseEntityInformation;
@@ -39,13 +35,16 @@
 import org.springframework.data.mapping.model.ConvertingPropertyAccessor;
 import org.springframework.util.ClassUtils;
 
-import com.couchbase.client.core.error.CouchbaseException;
+import java.lang.reflect.InaccessibleObjectException;
+import java.util.Map;
+import java.util.Set;
 
 
 /**
  * Base shared by Reactive and non-Reactive TemplateSupport
  *
  * @author Michael Reiche
+ * @author Mico Piira
  */
 @Stability.Internal
 public abstract class AbstractTemplateSupport {
@@ -68,7 +67,7 @@ public AbstractTemplateSupport(ReactiveCouchbaseTemplate template, CouchbaseConv
 	abstract ReactiveCouchbaseTemplate getReactiveTemplate();
 
 	public <T> T decodeEntityBase(Object id, String source, Long cas, Class<T> entityClass, String scope,
-			String collection, Object txResultHolder, CouchbaseResourceHolder holder) {
+			String collection, Object txResultHolder, CouchbaseResourceHolder holder, CouchbaseDocument converted) {
 
 		// this is the entity class defined for the repository. It may not be the class of the document that was read
 		// we will reset it after reading the document
@@ -88,7 +87,6 @@ public <T> T decodeEntityBase(Object id, String source, Long cas, Class<T> entit
 			// to unwrap. This results in List<String[]> being unwrapped past String[] to String, so this may also be a
 			// Collection (or Array) of entityClass. We have no way of knowing - so just assume it is what we are told.
 			// if this is a Collection or array, only the first element will be returned.
-			final CouchbaseDocument converted = new CouchbaseDocument(id);
 			Set<Map.Entry<String, Object>> set = ((CouchbaseDocument) translationService.decode(source, converted))
 					.getContent().entrySet();
 			return (T) set.iterator().next().getValue();
@@ -99,8 +97,6 @@ public <T> T decodeEntityBase(Object id, String source, Long cas, Class<T> entit
 					+ TemplateUtils.SELECT_ID);
 		}
 
-		final CouchbaseDocument converted = new CouchbaseDocument(id);
-
 		// if possible, set the version property in the source so that if the constructor has a long version argument,
 		// it will have a value and not fail (as null is not a valid argument for a long argument). This possible failure
 		// can be avoid by defining the argument as Long instead of long.
@@ -148,7 +144,7 @@ CouchbasePersistentEntity couldBePersistentEntity(Class<?> entityClass) {
 		return null;
 	}
 
-	public <T> T applyResultBase(T entity, CouchbaseDocument converted, Object id, long cas,
+	public <T> T applyResultBase(T entity, Object id, long cas,
 			Object txResultHolder, CouchbaseResourceHolder holder) {
 		ConvertingPropertyAccessor<Object> accessor = getPropertyAccessor(entity);
 
@@ -168,7 +164,6 @@ public <T> T applyResultBase(T entity, CouchbaseDocument converted, Object id, l
 		if (holder != null) {
 			holder.transactionResultHolder(txResultHolder, (T) accessor.getBean());
 		}
-		maybeEmitEvent(new AfterSaveEvent(accessor.getBean(), converted));
 		return (T) accessor.getBean();
 
 	}

src/main/java/org/springframework/data/couchbase/core/CouchbaseTemplateSupport.java

@@ -22,10 +22,7 @@
 import org.springframework.data.couchbase.core.convert.CouchbaseConverter;
 import org.springframework.data.couchbase.core.convert.translation.TranslationService;
 import org.springframework.data.couchbase.core.mapping.CouchbaseDocument;
-import org.springframework.data.couchbase.core.mapping.event.AfterConvertCallback;
-import org.springframework.data.couchbase.core.mapping.event.BeforeConvertCallback;
-import org.springframework.data.couchbase.core.mapping.event.BeforeConvertEvent;
-import org.springframework.data.couchbase.core.mapping.event.BeforeSaveEvent;
+import org.springframework.data.couchbase.core.mapping.event.*;
 import org.springframework.data.couchbase.transaction.CouchbaseResourceHolder;
 import org.springframework.data.mapping.callback.EntityCallbacks;
 import org.springframework.util.Assert;
@@ -37,6 +34,7 @@
  * @author Michael Reiche
  * @author Jorge Rodriguez Martin
  * @author Carlos Espinaco
+ * @author Mico Piira
  * @since 3.0
  */
 class CouchbaseTemplateSupport extends AbstractTemplateSupport implements ApplicationContextAware, TemplateSupport {
@@ -51,26 +49,30 @@ public CouchbaseTemplateSupport(final CouchbaseTemplate template, final Couchbas
 	}
 
 	@Override
-	public CouchbaseDocument encodeEntity(final Object entityToEncode) {
+	public <T> EncodedEntity<T> encodeEntity(final T entityToEncode) {
 		maybeEmitEvent(new BeforeConvertEvent<>(entityToEncode));
 		Object maybeNewEntity = maybeCallBeforeConvert(entityToEncode, "");
 		final CouchbaseDocument converted = new CouchbaseDocument();
 		converter.write(maybeNewEntity, converted);
-		maybeCallAfterConvert(entityToEncode, converted, "");
 		maybeEmitEvent(new BeforeSaveEvent<>(entityToEncode, converted));
-		return converted;
+		return new EncodedEntity<>(maybeCallBeforeSave(entityToEncode, converted, ""), converted);
 	}
 
 	@Override
 	public <T> T decodeEntity(Object id, String source, Long cas, Class<T> entityClass, String scope, String collection,
 			Object txHolder, CouchbaseResourceHolder holder) {
-		return decodeEntityBase(id, source, cas, entityClass, scope, collection, txHolder, holder);
+		CouchbaseDocument converted = new CouchbaseDocument(id);
+		T decoded = decodeEntityBase(id, source, cas, entityClass, scope, collection, txHolder, holder, converted);
+		maybeEmitEvent(new AfterConvertEvent<>(decoded, converted));
+		return maybeCallAfterConvert(decoded, converted, "");
 	}
 
 	@Override
 	public <T> T applyResult(T entity, CouchbaseDocument converted, Object id, long cas,
 			Object txResultHolder, CouchbaseResourceHolder holder) {
-		return applyResultBase(entity, converted, id, cas, txResultHolder, holder);
+		T applied = applyResultBase(entity, id, cas, txResultHolder, holder);
+		maybeEmitEvent(new AfterSaveEvent<>(applied, converted));
+		return maybeCallAfterSave(applied, converted, "");
 	}
 
 	@Override
@@ -119,6 +121,24 @@ protected <T> T maybeCallAfterConvert(T object, CouchbaseDocument document, Stri
 		return object;
 	}
 
+	protected <T> T maybeCallAfterSave(T object, CouchbaseDocument document, String collection) {
+		if (null != entityCallbacks) {
+			return entityCallbacks.callback(AfterSaveCallback.class, object, document, collection);
+		} else {
+			LOG.info("maybeCallAfterSave called, but CouchbaseTemplate not initialized with applicationContext");
+		}
+		return object;
+	}
+
+	protected <T> T maybeCallBeforeSave(T object, CouchbaseDocument document, String collection) {
+		if (null != entityCallbacks) {
+			return entityCallbacks.callback(BeforeSaveCallback.class, object, document, collection);
+		} else {
+			LOG.info("maybeCallBeforeSave called, but CouchbaseTemplate not initialized with applicationContext");
+		}
+		return object;
+	}
+
 	@Override
 	ReactiveCouchbaseTemplate getReactiveTemplate() {
 		return template.reactive();

src/main/java/org/springframework/data/couchbase/core/EncodedEntity.java

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2012-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.couchbase.core;
+
+import org.springframework.data.couchbase.core.mapping.CouchbaseDocument;
+
+/**
+ * @author Mico Piira
+ */
+public record EncodedEntity<T>(T entity, CouchbaseDocument document) {
+}

src/main/java/org/springframework/data/couchbase/core/NonReactiveSupportWrapper.java

@@ -26,6 +26,7 @@
  *
  * @author Carlos Espinaco
  * @author Michael Reiche
+ * @author Mico Piira
  * @since 4.2
  */
 public class NonReactiveSupportWrapper implements ReactiveTemplateSupport {
@@ -37,7 +38,7 @@ public NonReactiveSupportWrapper(TemplateSupport support) {
 	}
 
 	@Override
-	public Mono<CouchbaseDocument> encodeEntity(Object entityToEncode) {
+	public <T> Mono<EncodedEntity<T>> encodeEntity(T entityToEncode) {
 		return Mono.fromSupplier(() -> support.encodeEntity(entityToEncode));
 	}
 

src/main/java/org/springframework/data/couchbase/core/ReactiveCouchbaseTemplateSupport.java

@@ -16,28 +16,25 @@
 
 package org.springframework.data.couchbase.core;
 
-import reactor.core.publisher.Mono;
-
 import org.springframework.beans.BeansException;
 import org.springframework.context.ApplicationContext;
 import org.springframework.context.ApplicationContextAware;
 import org.springframework.data.couchbase.core.convert.CouchbaseConverter;
 import org.springframework.data.couchbase.core.convert.translation.TranslationService;
 import org.springframework.data.couchbase.core.mapping.CouchbaseDocument;
-import org.springframework.data.couchbase.core.mapping.event.BeforeConvertEvent;
-import org.springframework.data.couchbase.core.mapping.event.BeforeSaveEvent;
-import org.springframework.data.couchbase.core.mapping.event.ReactiveAfterConvertCallback;
-import org.springframework.data.couchbase.core.mapping.event.ReactiveBeforeConvertCallback;
+import org.springframework.data.couchbase.core.mapping.event.*;
 import org.springframework.data.couchbase.transaction.CouchbaseResourceHolder;
 import org.springframework.data.mapping.callback.EntityCallbacks;
 import org.springframework.data.mapping.callback.ReactiveEntityCallbacks;
 import org.springframework.util.Assert;
+import reactor.core.publisher.Mono;
 
 /**
  * Internal encode/decode support for {@link ReactiveCouchbaseTemplate}.
  *
  * @author Carlos Espinaco
  * @author Michael Reiche
+ * @author Mico Piira
  * @since 4.2
  */
 class ReactiveCouchbaseTemplateSupport extends AbstractTemplateSupport
@@ -53,14 +50,19 @@ public ReactiveCouchbaseTemplateSupport(final ReactiveCouchbaseTemplate template
 	}
 
 	@Override
-	public Mono<CouchbaseDocument> encodeEntity(final Object entityToEncode) {
-		return Mono.just(entityToEncode).doOnNext(entity -> maybeEmitEvent(new BeforeConvertEvent<>(entity)))
-				.flatMap(entity -> maybeCallBeforeConvert(entity, "")).map(maybeNewEntity -> {
+	public <T> Mono<EncodedEntity<T>> encodeEntity(final T entityToEncode) {
+		maybeEmitEvent(new BeforeConvertEvent<>(entityToEncode));
+		return maybeCallBeforeConvert(entityToEncode, "")
+				.map(maybeNewEntity -> {
 					final CouchbaseDocument converted = new CouchbaseDocument();
 					converter.write(maybeNewEntity, converted);
 					return converted;
-				}).flatMap(converted -> maybeCallAfterConvert(entityToEncode, converted, "").thenReturn(converted))
-				.doOnNext(converted -> maybeEmitEvent(new BeforeSaveEvent<>(entityToEncode, converted)));
+				})
+				.flatMap(converted -> {
+					maybeEmitEvent(new BeforeSaveEvent<>(entityToEncode, converted));
+					return maybeCallBeforeSave(entityToEncode, converted, "")
+							.map(potentiallyModified -> new EncodedEntity<>(potentiallyModified, converted));
+				});
 	}
 
 	@Override
@@ -71,14 +73,23 @@ ReactiveCouchbaseTemplate getReactiveTemplate() {
 	@Override
 	public <T> Mono<T> decodeEntity(Object id, String source, Long cas, Class<T> entityClass, String scope,
 			String collection, Object txResultHolder, CouchbaseResourceHolder holder) {
+		CouchbaseDocument converted = new CouchbaseDocument(id);
 		return Mono
-				.fromSupplier(() -> decodeEntityBase(id, source, cas, entityClass, scope, collection, txResultHolder, holder));
+				.fromSupplier(() -> decodeEntityBase(id, source, cas, entityClass, scope, collection, txResultHolder, holder, converted))
+				.flatMap(entity -> {
+					maybeEmitEvent(new AfterConvertEvent<>(entity, converted));
+					return maybeCallAfterConvert(entity, converted, "");
+				});
 	}
 
 	@Override
 	public <T> Mono<T> applyResult(T entity, CouchbaseDocument converted, Object id, Long cas,
 			Object txResultHolder, CouchbaseResourceHolder holder) {
-		return Mono.fromSupplier(() -> applyResultBase(entity, converted, id, cas, txResultHolder, holder));
+		return Mono.fromSupplier(() -> applyResultBase(entity, id, cas, txResultHolder, holder))
+				.flatMap(saved -> {
+					maybeEmitEvent(new AfterSaveEvent<>(saved, converted));
+					return maybeCallAfterSave(saved, converted, "");
+				});
 	}
 
 	@Override
@@ -104,6 +115,25 @@ public void setReactiveEntityCallbacks(ReactiveEntityCallbacks reactiveEntityCal
 		this.reactiveEntityCallbacks = reactiveEntityCallbacks;
 	}
 
+	protected <T> Mono<T> maybeCallBeforeSave(T object, CouchbaseDocument document, String collection) {
+		if (reactiveEntityCallbacks != null) {
+			return reactiveEntityCallbacks.callback(ReactiveBeforeSaveCallback.class, object, document, collection);
+		} else {
+			LOG.info("maybeCallBeforeSave called, but ReactiveCouchbaseTemplate not initialized with applicationContext");
+		}
+		return Mono.just(object);
+	}
+
+	protected <T> Mono<T> maybeCallAfterSave(T object, CouchbaseDocument document, String collection) {
+		if (reactiveEntityCallbacks != null) {
+			return reactiveEntityCallbacks.callback(ReactiveAfterSaveCallback.class, object, document, collection);
+		} else {
+			LOG.info("maybeCallAfterSave called, but ReactiveCouchbaseTemplate not initialized with applicationContext");
+		}
+		return Mono.just(object);
+	}
+
+
 	protected <T> Mono<T> maybeCallBeforeConvert(T object, String collection) {
 		if (reactiveEntityCallbacks != null) {
 			return reactiveEntityCallbacks.callback(ReactiveBeforeConvertCallback.class, object, collection);

src/main/java/org/springframework/data/couchbase/core/ReactiveInsertByIdOperationSupport.java

@@ -46,6 +46,7 @@
  *
  * @author Michael Reiche
  * @author Tigran Babloyan
+ * @author Mico Piira
  */
 public class ReactiveInsertByIdOperationSupport implements ReactiveInsertByIdOperation {
 
@@ -102,12 +103,15 @@ public Mono<T> one(T object) {
 			return Mono
 					.just(template.getCouchbaseClientFactory().withScope(pArgs.getScope()).getCollection(pArgs.getCollection()))
 					.flatMap(collection -> support.encodeEntity(object)
-							.flatMap(converted -> TransactionalSupport.checkForTransactionInThreadLocalStorage().flatMap(ctxOpt -> {
+							.flatMap(encodedEntity -> TransactionalSupport.checkForTransactionInThreadLocalStorage().flatMap(ctxOpt -> {
+								T potentiallyModified = encodedEntity.entity();
+								CouchbaseDocument converted = encodedEntity.document();
+
 								if (!ctxOpt.isPresent()) {
 									return collection.reactive()
 											.insert(converted.getId().toString(), converted.export(),
 													buildOptions(pArgs.getOptions(), converted))
-											.flatMap(result -> this.support.applyResult(object, converted, converted.getId(), result.cas(),
+											.flatMap(result -> this.support.applyResult(potentiallyModified, converted, converted.getId(), result.cas(),
 													null, null));
 								} else {
 									rejectInvalidTransactionalOptions();
@@ -120,7 +124,7 @@ public Mono<T> one(T object) {
 													template.getCouchbaseClientFactory().getCluster().environment().transcoder()
 															.encode(converted.export()).encoded(),
 													new SpanWrapper(span))
-											.flatMap(result -> this.support.applyResult(object, converted, converted.getId(), result.cas(),
+											.flatMap(result -> this.support.applyResult(potentiallyModified, converted, converted.getId(), result.cas(),
 													null, null));
 								}
 							})).onErrorMap(throwable -> {

src/main/java/org/springframework/data/couchbase/core/ReactiveMutateInByIdOperationSupport.java

@@ -23,7 +23,6 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.data.couchbase.core.mapping.CouchbaseDocument;
-import org.springframework.data.couchbase.core.mapping.CouchbaseList;
 import org.springframework.data.couchbase.core.query.OptionsBuilder;
 import org.springframework.data.couchbase.core.support.PseudoArgs;
 import org.springframework.util.Assert;
@@ -37,6 +36,7 @@
  * {@link ReactiveMutateInByIdOperation} implementations for Couchbase.
  *
  * @author Tigran Babloyan
+ * @author Mico Piira
  */
 public class ReactiveMutateInByIdOperationSupport implements ReactiveMutateInByIdOperation {
 
@@ -105,14 +105,16 @@ public Mono<T> one(T object) {
 			}
 
 			Mono<T> reactiveEntity = TransactionalSupport.verifyNotInTransaction("mutateInById")
-					.then(support.encodeEntity(object)).flatMap(converted -> {
+					.then(support.encodeEntity(object)).flatMap(encodedEntity -> {
+						T potentiallyModified = encodedEntity.entity();
+						CouchbaseDocument converted = encodedEntity.document();
 						return Mono
 								.just(template.getCouchbaseClientFactory().withScope(pArgs.getScope())
 										.getCollection(pArgs.getCollection()))
 								.flatMap(collection -> collection.reactive()
-										.mutateIn(converted.getId().toString(), getMutations(converted), buildMutateInOptions(pArgs.getOptions(), object, converted))
+										.mutateIn(converted.getId().toString(), getMutations(converted), buildMutateInOptions(pArgs.getOptions(), potentiallyModified, converted))
 										.flatMap(
-												result -> support.applyResult(object, converted, converted.getId(), result.cas(), null, null)));
+												result -> support.applyResult(potentiallyModified, converted, converted.getId(), result.cas(), null, null)));
 					});
 
 			return reactiveEntity.onErrorMap(throwable -> {