Capture execution results for update db actions to fix root references with immutably typed ids.

+ Added #getIdValueSource to DbAction.WithEntity since it now applies to all such DbAction.
+ Removed DbAction.WithGeneratedId in favor of DbAction.WithEntity where it was used.

Removes unused DbAction.Update.

Passes the root from BeforeSaveCallback through to InsertRoot,UpdateRoot
actions.

+ Isolate aggregate change #getEntity #setEntity behaviors into a new
interface for "save" changes.
+ Extract #setEntity from DbAction.WithEntity to a new DbAction interface
just for root actions.
+ JdbcAggregateChangeExecutionContext#populateIdsIfNecessary returns the root entity
from DbAction.InsertRoot or DbAction.UpdateRoot regardless of the immutability of
its @id property.

Closes #1199
Closes #1201
Original pull request #1208

Author: ctailor2

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/AggregateChangeExecutor.java

@@ -18,10 +18,10 @@
 import org.springframework.data.jdbc.core.convert.DataAccessStrategy;
 import org.springframework.data.jdbc.core.convert.JdbcConverter;
 import org.springframework.data.relational.core.conversion.AggregateChange;
+import org.springframework.data.relational.core.conversion.AggregateChangeWithRoot;
 import org.springframework.data.relational.core.conversion.DbAction;
 import org.springframework.data.relational.core.conversion.DbActionExecutionException;
 import org.springframework.data.relational.core.conversion.MutableAggregateChange;
-import org.springframework.lang.Nullable;
 
 /**
  * Executes an {@link MutableAggregateChange}.
@@ -42,20 +42,22 @@ class AggregateChangeExecutor {
 		this.accessStrategy = accessStrategy;
 	}
 
-	@Nullable
-	<T> T execute(AggregateChange<T> aggregateChange) {
+	<T> T execute(AggregateChangeWithRoot<T> aggregateChange) {
 
 		JdbcAggregateChangeExecutionContext executionContext = new JdbcAggregateChangeExecutionContext(converter,
 				accessStrategy);
 
 		aggregateChange.forEachAction(action -> execute(action, executionContext));
 
-		T root = executionContext.populateIdsIfNecessary();
-		if (root == null) {
-			root = aggregateChange.getEntity();
-		}
+		return executionContext.populateIdsIfNecessary();
+	}
 
-		return root;
+	<T> void execute(AggregateChange<T> aggregateChange) {
+
+		JdbcAggregateChangeExecutionContext executionContext = new JdbcAggregateChangeExecutionContext(converter,
+				accessStrategy);
+
+		aggregateChange.forEachAction(action -> execute(action, executionContext));
 	}
 
 	private void execute(DbAction<?> action, JdbcAggregateChangeExecutionContext executionContext) {
@@ -69,8 +71,6 @@ private void execute(DbAction<?> action, JdbcAggregateChangeExecutionContext exe
 				executionContext.executeInsertBatch((DbAction.InsertBatch<?>) action);
 			} else if (action instanceof DbAction.UpdateRoot) {
 				executionContext.executeUpdateRoot((DbAction.UpdateRoot<?>) action);
-			} else if (action instanceof DbAction.Update) {
-				executionContext.executeUpdate((DbAction.Update<?>) action);
 			} else if (action instanceof DbAction.Delete) {
 				executionContext.executeDelete((DbAction.Delete<?>) action);
 			} else if (action instanceof DbAction.DeleteAll) {

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutionContext.java

@@ -15,15 +15,7 @@
  */
 package org.springframework.data.jdbc.core;
 
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.LinkedHashMap;
-import java.util.List;
-import java.util.Map;
-import java.util.Set;
+import java.util.*;
 import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
@@ -40,6 +32,7 @@
 import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.relational.core.conversion.DbAction;
 import org.springframework.data.relational.core.conversion.DbActionExecutionResult;
+import org.springframework.data.relational.core.conversion.IdValueSource;
 import org.springframework.data.relational.core.mapping.PersistentPropertyPathExtension;
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
@@ -108,15 +101,7 @@ <T> void executeUpdateRoot(DbAction.UpdateRoot<T> update) {
 		} else {
 			updateWithoutVersion(update);
 		}
-	}
-
-	<T> void executeUpdate(DbAction.Update<T> update) {
-
-		if (!accessStrategy.update(update.getEntity(), update.getEntityType())) {
-
-			throw new IncorrectUpdateSemanticsDataAccessException(
-					String.format(UPDATE_FAILED, update.getEntity(), getIdFrom(update)));
-		}
+		add(new DbActionExecutionResult(update));
 	}
 
 	<T> void executeDeleteRoot(DbAction.DeleteRoot<T> delete) {
@@ -203,11 +188,12 @@ private DbAction.WithEntity<?> getIdOwningAction(DbAction.WithEntity<?> action,
 
 	private Object getPotentialGeneratedIdFrom(DbAction.WithEntity<?> idOwningAction) {
 
-		if (idOwningAction instanceof DbAction.WithGeneratedId) {
+		if (IdValueSource.GENERATED.equals(idOwningAction.getIdValueSource())) {
 
-			Object generatedId;
 			DbActionExecutionResult dbActionExecutionResult = results.get(idOwningAction);
-			generatedId = dbActionExecutionResult == null ? null : dbActionExecutionResult.getId();
+			Object generatedId = Optional.ofNullable(dbActionExecutionResult) //
+					.map(DbActionExecutionResult::getGeneratedId) //
+					.orElse(null);
 
 			if (generatedId != null) {
 				return generatedId;
@@ -227,12 +213,8 @@ private Object getIdFrom(DbAction.WithEntity<?> idOwningAction) {
 		return identifier;
 	}
 
-	@SuppressWarnings("unchecked")
-	@Nullable
 	<T> T populateIdsIfNecessary() {
 
-		T newRoot = null;
-
 		// have the results so that the inserts on the leaves come first.
 		List<DbActionExecutionResult> reverseResults = new ArrayList<>(results.values());
 		Collections.reverse(reverseResults);
@@ -241,33 +223,29 @@ <T> T populateIdsIfNecessary() {
 
 		for (DbActionExecutionResult result : reverseResults) {
 
-			DbAction<?> action = result.getAction();
+			DbAction.WithEntity<?> action = result.getAction();
 
-			if (!(action instanceof DbAction.WithGeneratedId)) {
-				continue;
+			Object newEntity = setIdAndCascadingProperties(action, result.getGeneratedId(), cascadingValues);
+			
+			if (action instanceof DbAction.InsertRoot || action instanceof DbAction.UpdateRoot) {
+				//noinspection unchecked
+				return (T) newEntity;
 			}
 
-			DbAction.WithEntity<?> withEntity = (DbAction.WithGeneratedId<?>) action;
-			Object newEntity = setIdAndCascadingProperties(withEntity, result.getId(), cascadingValues);
-
 			// the id property was immutable so we have to propagate changes up the tree
-			if (newEntity != withEntity.getEntity()) {
-
-				if (action instanceof DbAction.Insert) {
-					DbAction.Insert<?> insert = (DbAction.Insert<?>) action;
-
-					Pair<?, ?> qualifier = insert.getQualifier();
+			if (newEntity != action.getEntity() && action instanceof DbAction.Insert) {
+				DbAction.Insert<?> insert = (DbAction.Insert<?>) action;
 
-					cascadingValues.stage(insert.getDependingOn(), insert.getPropertyPath(),
-							qualifier == null ? null : qualifier.getSecond(), newEntity);
+				Pair<?, ?> qualifier = insert.getQualifier();
 
-				} else if (action instanceof DbAction.InsertRoot) {
-					newRoot = (T) newEntity;
-				}
+				cascadingValues.stage(insert.getDependingOn(), insert.getPropertyPath(),
+						qualifier == null ? null : qualifier.getSecond(), newEntity);
 			}
 		}
 
-		return newRoot;
+		throw new IllegalStateException(
+				String.format("Cannot retrieve the resulting instance unless a %s or %s action was successfully executed.",
+						DbAction.InsertRoot.class.getName(), DbAction.UpdateRoot.class.getName()));
 	}
 
 	private <S> Object setIdAndCascadingProperties(DbAction.WithEntity<S> action, @Nullable Object generatedId,
@@ -279,7 +257,7 @@ private <S> Object setIdAndCascadingProperties(DbAction.WithEntity<S> action, @N
 				.getRequiredPersistentEntity(action.getEntityType());
 		PersistentPropertyAccessor<S> propertyAccessor = converter.getPropertyAccessor(persistentEntity, originalEntity);
 
-		if (generatedId != null && persistentEntity.hasIdProperty()) {
+		if (IdValueSource.GENERATED.equals(action.getIdValueSource())) {
 			propertyAccessor.setProperty(persistentEntity.getRequiredIdProperty(), generatedId);
 		}
 
@@ -301,6 +279,10 @@ private PersistentPropertyPath<?> getRelativePath(DbAction<?> action, Persistent
 			return pathToValue;
 		}
 
+		if (action instanceof DbAction.UpdateRoot) {
+			return pathToValue;
+		}
+
 		throw new IllegalArgumentException(String.format("DbAction of type %s is not supported.", action.getClass()));
 	}
 

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateTemplate.java

@@ -31,6 +31,7 @@
 import org.springframework.data.mapping.IdentifierAccessor;
 import org.springframework.data.mapping.callback.EntityCallbacks;
 import org.springframework.data.relational.core.conversion.AggregateChange;
+import org.springframework.data.relational.core.conversion.AggregateChangeWithRoot;
 import org.springframework.data.relational.core.conversion.MutableAggregateChange;
 import org.springframework.data.relational.core.conversion.RelationalEntityDeleteWriter;
 import org.springframework.data.relational.core.conversion.RelationalEntityInsertWriter;
@@ -61,8 +62,6 @@ public class JdbcAggregateTemplate implements JdbcAggregateOperations {
 	private final RelationalMappingContext context;
 
 	private final RelationalEntityDeleteWriter jdbcEntityDeleteWriter;
-	private final RelationalEntityInsertWriter jdbcEntityInsertWriter;
-	private final RelationalEntityUpdateWriter jdbcEntityUpdateWriter;
 
 	private final DataAccessStrategy accessStrategy;
 	private final AggregateChangeExecutor executor;
@@ -92,8 +91,6 @@ public JdbcAggregateTemplate(ApplicationContext publisher, RelationalMappingCont
 		this.accessStrategy = dataAccessStrategy;
 		this.converter = converter;
 
-		this.jdbcEntityInsertWriter = new RelationalEntityInsertWriter(context);
-		this.jdbcEntityUpdateWriter = new RelationalEntityUpdateWriter(context);
 		this.jdbcEntityDeleteWriter = new RelationalEntityDeleteWriter(context);
 
 		this.executor = new AggregateChangeExecutor(converter, accessStrategy);
@@ -122,8 +119,6 @@ public JdbcAggregateTemplate(ApplicationEventPublisher publisher, RelationalMapp
 		this.accessStrategy = dataAccessStrategy;
 		this.converter = converter;
 
-		this.jdbcEntityInsertWriter = new RelationalEntityInsertWriter(context);
-		this.jdbcEntityUpdateWriter = new RelationalEntityUpdateWriter(context);
 		this.jdbcEntityDeleteWriter = new RelationalEntityDeleteWriter(context);
 		this.executor = new AggregateChangeExecutor(converter, accessStrategy);
 	}
@@ -146,7 +141,7 @@ public <T> T save(T instance) {
 
 		RelationalPersistentEntity<?> persistentEntity = context.getRequiredPersistentEntity(instance.getClass());
 
-		Function<T, MutableAggregateChange<T>> changeCreator = persistentEntity.isNew(instance)
+		Function<T, AggregateChangeWithRoot<T>> changeCreator = persistentEntity.isNew(instance)
 				? entity -> createInsertChange(prepareVersionForInsert(entity))
 				: entity -> createUpdateChange(prepareVersionForUpdate(entity));
 
@@ -286,18 +281,18 @@ public void deleteAll(Class<?> domainType) {
 		executor.execute(change);
 	}
 
-	private <T> T store(T aggregateRoot, Function<T, MutableAggregateChange<T>> changeCreator,
+	private <T> T store(T aggregateRoot, Function<T, AggregateChangeWithRoot<T>> changeCreator,
 			RelationalPersistentEntity<?> persistentEntity) {
 
 		Assert.notNull(aggregateRoot, "Aggregate instance must not be null!");
 
 		aggregateRoot = triggerBeforeConvert(aggregateRoot);
 
-		MutableAggregateChange<T> change = changeCreator.apply(aggregateRoot);
+		AggregateChangeWithRoot<T> change = changeCreator.apply(aggregateRoot);
 
-		aggregateRoot = triggerBeforeSave(change.getEntity(), change);
+		aggregateRoot = triggerBeforeSave(change.getRoot(), change);
 
-		change.setEntity(aggregateRoot);
+		change.setRoot(aggregateRoot);
 
 		T entityAfterExecution = executor.execute(change);
 
@@ -313,25 +308,24 @@ private <T> void deleteTree(Object id, @Nullable T entity, Class<T> domainType)
 		MutableAggregateChange<T> change = createDeletingChange(id, entity, domainType);
 
 		entity = triggerBeforeDelete(entity, id, change);
-		change.setEntity(entity);
 
 		executor.execute(change);
 
 		triggerAfterDelete(entity, id, change);
 	}
 
-	private <T> MutableAggregateChange<T> createInsertChange(T instance) {
+	private <T> AggregateChangeWithRoot<T> createInsertChange(T instance) {
 
-		MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forSave(instance);
-		jdbcEntityInsertWriter.write(instance, aggregateChange);
+		AggregateChangeWithRoot<T> aggregateChange = MutableAggregateChange.forSave(instance);
+		new RelationalEntityInsertWriter<T>(context).write(instance, aggregateChange);
 		return aggregateChange;
 	}
 
-	private <T> MutableAggregateChange<T> createUpdateChange(EntityAndPreviousVersion<T> entityAndVersion) {
+	private <T> AggregateChangeWithRoot<T> createUpdateChange(EntityAndPreviousVersion<T> entityAndVersion) {
 
-		MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forSave(entityAndVersion.entity,
+		AggregateChangeWithRoot<T> aggregateChange = MutableAggregateChange.forSave(entityAndVersion.entity,
 				entityAndVersion.version);
-		jdbcEntityUpdateWriter.write(entityAndVersion.entity, aggregateChange);
+		new RelationalEntityUpdateWriter<T>(context).write(entityAndVersion.entity, aggregateChange);
 		return aggregateChange;
 	}
 
@@ -383,14 +377,14 @@ private <T> MutableAggregateChange<T> createDeletingChange(Object id, @Nullable
 				previousVersion = RelationalEntityVersionUtils.getVersionNumberFromEntity(entity, persistentEntity, converter);
 			}
 		}
-		MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forDelete(domainType, entity, previousVersion);
+		MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forDelete(domainType, previousVersion);
 		jdbcEntityDeleteWriter.write(id, aggregateChange);
 		return aggregateChange;
 	}
 
 	private MutableAggregateChange<?> createDeletingChange(Class<?> domainType) {
 
-		MutableAggregateChange<?> aggregateChange = MutableAggregateChange.forDelete(domainType, null);
+		MutableAggregateChange<?> aggregateChange = MutableAggregateChange.forDelete(domainType);
 		jdbcEntityDeleteWriter.write(null, aggregateChange);
 		return aggregateChange;
 	}