DATAMONGO-2307 - Consistently use result of BeforeSaveCallback.

Original Pull Request: #765

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/MongoTemplate.java

@@ -169,6 +169,7 @@ import com.mongodb.client.result.UpdateResult;
  * @author duozhilin
  * @author Andreas Zink
  * @author Cimon Lucas
+ * @author Michael J. Simons
  */
 @SuppressWarnings("deprecation")
 public class MongoTemplate implements MongoOperations, ApplicationContextAware, IndexOperationsProvider {
@@ -1466,7 +1467,7 @@ public class MongoTemplate implements MongoOperations, ApplicationContextAware,
 		objectToSave = maybeCallBeforeSave(objectToSave, dbDoc, collectionName);
 		Object id = saveDocument(collectionName, dbDoc, objectToSave.getClass());
 
-		T saved = populateIdIfNecessary(entity.getBean(), id);
+		T saved = populateIdIfNecessary(objectToSave, id);
 		maybeEmitEvent(new AfterSaveEvent<>(saved, dbDoc, collectionName));
 
 		return saved;

spring-data-mongodb/src/main/java/org/springframework/data/mongodb/core/mapping/event/BeforeSaveCallback.java

@@ -22,6 +22,7 @@ import org.springframework.data.mapping.callback.EntityCallback;
  * Entity callback triggered before save of a document.
  *
  * @author Mark Paluch
+ * @author Michael J. Simons
  * @since 2.2
  */
 @FunctionalInterface
@@ -29,8 +30,11 @@ public interface BeforeSaveCallback<T> extends EntityCallback<T> {
 
 	/**
 	 * Entity callback method invoked before a domain object is saved. Can return either the same or a modified instance
-	 * of the domain object and can modify {@link Document} contents. This method called after converting the
-	 * {@code entity} to {@link Document} so effectively the document is used as outcome of invoking this callback.
+	 * of the domain object and can modify {@link Document} contents. This method is called after converting the
+	 * {@code entity} to a {@link Document} so effectively the document is used as outcome of invoking this callback.
+	 * Changes to the domain object are not taken into account for saving, only changes to the document. Only transient
+	 * fields of the entity should be changed in this callback. To change persistent the entity before being converted,
+	 * use the {@link BeforeConvertCallback}.
 	 *
 	 * @param entity the domain object to save.
 	 * @param document {@link Document} representing the {@code entity}.

spring-data-mongodb/src/test/java/org/springframework/data/mongodb/core/MongoTemplateUnitTests.java

@@ -53,6 +53,7 @@ import org.springframework.core.convert.converter.Converter;
 import org.springframework.dao.DataAccessException;
 import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.data.annotation.Id;
+import org.springframework.data.annotation.Transient;
 import org.springframework.data.annotation.Version;
 import org.springframework.data.convert.CustomConversions;
 import org.springframework.data.domain.Sort;
@@ -84,6 +85,7 @@ import org.springframework.data.mongodb.core.query.Query;
 import org.springframework.data.mongodb.core.query.Update;
 import org.springframework.lang.Nullable;
 import org.springframework.test.util.ReflectionTestUtils;
+import org.springframework.transaction.annotation.Transactional;
 import org.springframework.util.CollectionUtils;
 
 import com.mongodb.DB;
@@ -117,6 +119,7 @@ import com.mongodb.client.result.UpdateResult;
  * @author Oliver Gierke
  * @author Christoph Strobl
  * @author Mark Paluch
+ * @author Michael J. Simons
  */
 @RunWith(MockitoJUnitRunner.class)
 public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
@@ -1558,6 +1561,42 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 		assertThat(captor.getValue()).containsEntry("added-by", "callback");
 	}
 
+	@Test // DATAMONGO-2307
+	public void beforeSaveCallbackAllowsTargetEntityModificationsUsingSave() {
+
+		StaticApplicationContext ctx = new StaticApplicationContext();
+		ctx.registerBean(BeforeSaveCallback.class, () -> markPersonAsPersistedCallback);
+		ctx.refresh();
+
+		template.setApplicationContext(ctx);
+
+		PersonWithTransientAttribute entity = new PersonWithTransientAttribute();
+		entity.id = "luke-skywalker";
+		entity.firstname = "luke";
+		entity.isNew = true;
+
+		PersonWithTransientAttribute savedPerson = template.save(entity);
+		assertThat(savedPerson.isNew).isFalse();
+	}
+
+	@Test // DATAMONGO-2307
+	public void beforeSaveCallbackAllowsTargetEntityModificationsUsingInsert() {
+
+		StaticApplicationContext ctx = new StaticApplicationContext();
+		ctx.registerBean(BeforeSaveCallback.class, () -> markPersonAsPersistedCallback);
+		ctx.refresh();
+
+		template.setApplicationContext(ctx);
+
+		PersonWithTransientAttribute entity = new PersonWithTransientAttribute();
+		entity.id = "luke-skywalker";
+		entity.firstname = "luke";
+		entity.isNew = true;
+
+		PersonWithTransientAttribute savedPerson = template.insert(entity);
+		assertThat(savedPerson.isNew).isFalse();
+	}
+
 	// TODO: additional tests for what is when saved.
 
 	@Test // DATAMONGO-2261
@@ -1640,6 +1679,22 @@ public class MongoTemplateUnitTests extends MongoOperationsUnitTests {
 		String lastname;
 	}
 
+	static class PersonWithTransientAttribute extends Person {
+
+		@Transient
+		boolean isNew = true;
+	}
+
+	static BeforeSaveCallback<PersonWithTransientAttribute> markPersonAsPersistedCallback = (entity, document, collection) -> {
+
+		// Return a completely new instance, ie in case of an immutable entity;
+		PersonWithTransientAttribute newEntity = new PersonWithTransientAttribute();
+		newEntity.id = entity.id;
+		newEntity.firstname = entity.firstname;
+		newEntity.isNew = false;
+		return newEntity;
+	};
+
 	interface PersonProjection {
 		String getFirstname();
 	}