Determine and set the value for entity @Version before conversion to DbActions to simplify execution context.

This change incorporates one test from #1150 Original pull request #1196 Closes #1137

Author: ctailor2

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

@@ -51,10 +51,8 @@ <T> T execute(AggregateChange<T> aggregateChange) {
 aggregateChange.forEachAction(action -> execute(action, executionContext));
 
 T root = executionContext.populateIdsIfNecessary();
-root = root == null ? aggregateChange.getEntity() : root;
-
-if (root != null) {
-root = executionContext.populateRootVersionIfNecessary(root);
+if (root == null) {
+root = aggregateChange.getEntity();
 }
 
 return root;

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

@@ -40,7 +40,6 @@
 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.RelationalEntityVersionUtils;
 import org.springframework.data.relational.core.mapping.PersistentPropertyPathExtension;
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
 import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
@@ -65,7 +64,6 @@ class JdbcAggregateChangeExecutionContext {
 private final DataAccessStrategy accessStrategy;
 
 private final Map<DbAction<?>, DbActionExecutionResult> results = new LinkedHashMap<>();
-@Nullable private Long version;
 
 JdbcAggregateChangeExecutionContext(JdbcConverter converter, DataAccessStrategy accessStrategy) {
 
@@ -76,28 +74,8 @@ class JdbcAggregateChangeExecutionContext {
 
 <T> void executeInsertRoot(DbAction.InsertRoot<T> insert) {
 
-RelationalPersistentEntity<T> persistentEntity = getRequiredPersistentEntity(insert.getEntityType());
-
-Object id;
-if (persistentEntity.hasVersionProperty()) {
-
-RelationalPersistentProperty versionProperty = persistentEntity.getVersionProperty();
-
-Assert.state(versionProperty != null, "Version property must not be null at this stage.");
-
-long initialVersion = versionProperty.getActualType().isPrimitive() ? 1L : 0;
-
-T rootEntity = RelationalEntityVersionUtils.setVersionNumberOnEntity( //
-insert.getEntity(), initialVersion, persistentEntity, converter);
-
-id = accessStrategy.insert(rootEntity, insert.getEntityType(), Identifier.empty(), insert.getIdValueSource());
-
-setNewVersion(initialVersion);
-} else {
-id = accessStrategy.insert(insert.getEntity(), insert.getEntityType(), Identifier.empty(),
-insert.getIdValueSource());
-}
-
+Object id = accessStrategy.insert(insert.getEntity(), insert.getEntityType(), Identifier.empty(),
+insert.getIdValueSource());
 add(new DbActionExecutionResult(insert, id));
 }
 
@@ -125,12 +103,9 @@ <T> void executeInsertBatch(DbAction.InsertBatch<T> insertBatch) {
 
 <T> void executeUpdateRoot(DbAction.UpdateRoot<T> update) {
 
-RelationalPersistentEntity<T> persistentEntity = getRequiredPersistentEntity(update.getEntityType());
-
-if (persistentEntity.hasVersionProperty()) {
-updateWithVersion(update, persistentEntity);
+if (update.getPreviousVersion() != null) {
+updateWithVersion(update);
 } else {
-
 updateWithoutVersion(update);
 }
 }
@@ -147,16 +122,10 @@ <T> void executeUpdate(DbAction.Update<T> update) {
 <T> void executeDeleteRoot(DbAction.DeleteRoot<T> delete) {
 
 if (delete.getPreviousVersion() != null) {
-
-RelationalPersistentEntity<T> persistentEntity = getRequiredPersistentEntity(delete.getEntityType());
-if (persistentEntity.hasVersionProperty()) {
-
-accessStrategy.deleteWithVersion(delete.getId(), delete.getEntityType(), delete.getPreviousVersion());
-return;
-}
+accessStrategy.deleteWithVersion(delete.getId(), delete.getEntityType(), delete.getPreviousVersion());
+} else {
+accessStrategy.delete(delete.getId(), delete.getEntityType());
 }
-
-accessStrategy.delete(delete.getId(), delete.getEntityType());
 }
 
 <T> void executeDelete(DbAction.Delete<T> delete) {
@@ -258,36 +227,6 @@ private Object getIdFrom(DbAction.WithEntity<?> idOwningAction) {
 return identifier;
 }
 
-private void setNewVersion(long version) {
-
-Assert.isNull(this.version, "A new version was set a second time.");
-
-this.version = version;
-}
-
-private long getNewVersion() {
-
-Assert.notNull(version, "A new version was requested, but none was set.");
-
-return version;
-}
-
-private boolean hasNewVersion() {
-return version != null;
-}
-
-<T> T populateRootVersionIfNecessary(T newRoot) {
-
-if (!hasNewVersion()) {
-return newRoot;
-}
-// Does the root entity have a version attribute?
-RelationalPersistentEntity<T> persistentEntity = (RelationalPersistentEntity<T>) context
-.getRequiredPersistentEntity(newRoot.getClass());
-
-return RelationalEntityVersionUtils.setVersionNumberOnEntity(newRoot, getNewVersion(), persistentEntity, converter);
-}
-
 @SuppressWarnings("unchecked")
 @Nullable
 <T> T populateIdsIfNecessary() {
@@ -378,20 +317,12 @@ private <T> void updateWithoutVersion(DbAction.UpdateRoot<T> update) {
 }
 }
 
-private <T> void updateWithVersion(DbAction.UpdateRoot<T> update, RelationalPersistentEntity<T> persistentEntity) {
-
-// If the root aggregate has a version property, increment it.
-Number previousVersion = RelationalEntityVersionUtils.getVersionNumberFromEntity(update.getEntity(),
-persistentEntity, converter);
+private <T> void updateWithVersion(DbAction.UpdateRoot<T> update) {
 
+Number previousVersion = update.getPreviousVersion();
 Assert.notNull(previousVersion, "The root aggregate cannot be updated because the version property is null.");
 
-setNewVersion(previousVersion.longValue() + 1);
-
-T rootEntity = RelationalEntityVersionUtils.setVersionNumberOnEntity(update.getEntity(), getNewVersion(),
-persistentEntity, converter);
-
-if (!accessStrategy.updateWithVersion(rootEntity, update.getEntityType(), previousVersion)) {
+if (!accessStrategy.updateWithVersion(update.getEntity(), update.getEntityType(), previousVersion)) {
 
 throw new OptimisticLockingFailureException(String.format(UPDATE_FAILED_OPTIMISTIC_LOCKING, update.getEntity()));
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright 2017-2021 the original author or authors.
+ * Copyright 2017-2022 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.
@@ -35,8 +35,10 @@
 import org.springframework.data.relational.core.conversion.RelationalEntityDeleteWriter;
 import org.springframework.data.relational.core.conversion.RelationalEntityInsertWriter;
 import org.springframework.data.relational.core.conversion.RelationalEntityUpdateWriter;
+import org.springframework.data.relational.core.conversion.RelationalEntityVersionUtils;
 import org.springframework.data.relational.core.mapping.RelationalMappingContext;
 import org.springframework.data.relational.core.mapping.RelationalPersistentEntity;
+import org.springframework.data.relational.core.mapping.RelationalPersistentProperty;
 import org.springframework.data.relational.core.mapping.event.*;
 import org.springframework.data.support.PageableExecutionUtils;
 import org.springframework.lang.Nullable;
@@ -51,6 +53,7 @@
  * @author Christoph Strobl
  * @author Milan Milanov
  * @author Myeonghyeon Lee
+ * @author Chirag Tailor
  */
 public class JdbcAggregateTemplate implements JdbcAggregateOperations {
 
@@ -63,6 +66,7 @@ public class JdbcAggregateTemplate implements JdbcAggregateOperations {
 
 private final DataAccessStrategy accessStrategy;
 private final AggregateChangeExecutor executor;
+private final JdbcConverter converter;
 
 private EntityCallbacks entityCallbacks = EntityCallbacks.create();
 
@@ -86,6 +90,7 @@ public JdbcAggregateTemplate(ApplicationContext publisher, RelationalMappingCont
 this.publisher = publisher;
 this.context = context;
 this.accessStrategy = dataAccessStrategy;
+this.converter = converter;
 
 this.jdbcEntityInsertWriter = new RelationalEntityInsertWriter(context);
 this.jdbcEntityUpdateWriter = new RelationalEntityUpdateWriter(context);
@@ -115,6 +120,7 @@ public JdbcAggregateTemplate(ApplicationEventPublisher publisher, RelationalMapp
 this.publisher = publisher;
 this.context = context;
 this.accessStrategy = dataAccessStrategy;
+this.converter = converter;
 
 this.jdbcEntityInsertWriter = new RelationalEntityInsertWriter(context);
 this.jdbcEntityUpdateWriter = new RelationalEntityUpdateWriter(context);
@@ -288,7 +294,7 @@ private <T> T store(T aggregateRoot, Function<T, MutableAggregateChange<T>> chan
 
 MutableAggregateChange<T> change = changeCreator.apply(aggregateRoot);
 
-aggregateRoot = triggerBeforeSave(aggregateRoot, change);
+aggregateRoot = triggerBeforeSave(change.getEntity(), change);
 
 change.setEntity(aggregateRoot);
 
@@ -315,21 +321,58 @@ private <T> void deleteTree(Object id, @Nullable T entity, Class<T> domainType)
 
 private <T> MutableAggregateChange<T> createInsertChange(T instance) {
 
-MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forSave(instance);
-jdbcEntityInsertWriter.write(instance, aggregateChange);
+RelationalPersistentEntity<T> persistentEntity = getRequiredPersistentEntity(instance);
+T preparedInstance = instance;
+if (persistentEntity.hasVersionProperty()) {
+RelationalPersistentProperty versionProperty = persistentEntity.getRequiredVersionProperty();
+
+long initialVersion = versionProperty.getActualType().isPrimitive() ? 1L : 0;
+
+preparedInstance = RelationalEntityVersionUtils.setVersionNumberOnEntity( //
+instance, initialVersion, persistentEntity, converter);
+}
+MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forSave(preparedInstance);
+jdbcEntityInsertWriter.write(preparedInstance, aggregateChange);
 return aggregateChange;
 }
 
 private <T> MutableAggregateChange<T> createUpdateChange(T instance) {
 
-MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forSave(instance);
-jdbcEntityUpdateWriter.write(instance, aggregateChange);
+RelationalPersistentEntity<T> persistentEntity = getRequiredPersistentEntity(instance);
+T preparedInstance = instance;
+Number previousVersion = null;
+if (persistentEntity.hasVersionProperty()) {
+// If the root aggregate has a version property, increment it.
+previousVersion = RelationalEntityVersionUtils.getVersionNumberFromEntity(instance,
+persistentEntity, converter);
+
+Assert.notNull(previousVersion, "The root aggregate cannot be updated because the version property is null.");
+
+long newVersion = previousVersion.longValue() + 1;
+
+preparedInstance = RelationalEntityVersionUtils.setVersionNumberOnEntity(instance, newVersion,
+persistentEntity, converter);
+}
+MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forSave(preparedInstance, previousVersion);
+jdbcEntityUpdateWriter.write(preparedInstance, aggregateChange);
 return aggregateChange;
 }
 
+@SuppressWarnings("unchecked")
+private <T> RelationalPersistentEntity<T> getRequiredPersistentEntity(T instance) {
+return (RelationalPersistentEntity<T>) context.getRequiredPersistentEntity(instance.getClass());
+}
+
 private <T> MutableAggregateChange<T> createDeletingChange(Object id, @Nullable T entity, Class<T> domainType) {
 
-MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forDelete(domainType, entity);
+Number previousVersion = null;
+if (entity != null) {
+RelationalPersistentEntity<T> persistentEntity = getRequiredPersistentEntity(entity);
+if (persistentEntity.hasVersionProperty()) {
+previousVersion = RelationalEntityVersionUtils.getVersionNumberFromEntity(entity, persistentEntity, converter);
+}
+}
+MutableAggregateChange<T> aggregateChange = MutableAggregateChange.forDelete(domainType, entity, previousVersion);
 jdbcEntityDeleteWriter.write(id, aggregateChange);
 return aggregateChange;
 }

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutorContextImmutableUnitTests.java

@@ -64,7 +64,7 @@ public void rootOfEmptySetOfActionsisNull() {
 }
 
 @Test // DATAJDBC-453
-public void afterInsertRootIdAndVersionMaybeUpdated() {
+public void afterInsertRootIdMaybeUpdated() {
 
 // note that the root entity isn't the original one, but a new instance with the version set.
 when(accessStrategy.insert(any(DummyEntity.class), eq(DummyEntity.class), eq(Identifier.empty()),
@@ -76,10 +76,6 @@ public void afterInsertRootIdAndVersionMaybeUpdated() {
 
 assertThat(newRoot).isNotNull();
 assertThat(newRoot.id).isEqualTo(23L);
-
-newRoot = executionContext.populateRootVersionIfNecessary(newRoot);
-
-assertThat(newRoot.version).isEqualTo(1);
 }
 
 @Test // DATAJDBC-453

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutorContextUnitTests.java

@@ -70,7 +70,7 @@ public void rootOfEmptySetOfActionIsNull() {
 }
 
 @Test // DATAJDBC-453
-public void afterInsertRootIdAndVersionMaybeUpdated() {
+public void afterInsertRootIdMaybeUpdated() {
 
 when(accessStrategy.insert(root, DummyEntity.class, Identifier.empty(), IdValueSource.GENERATED)).thenReturn(23L);
 
@@ -80,22 +80,6 @@ public void afterInsertRootIdAndVersionMaybeUpdated() {
 
 assertThat(newRoot).isNull();
 assertThat(root.id).isEqualTo(23L);
-
-executionContext.populateRootVersionIfNecessary(root);
-
-assertThat(root.version).isEqualTo(1);
-}
-
-@Test // DATAJDBC-507
-public void afterInsertNotPrimitiveVersionShouldBeZero() {
-
-DummyEntityNonPrimitiveVersion dummyEntityNonPrimitiveVersion = new DummyEntityNonPrimitiveVersion();
-
-executionContext
-.executeInsertRoot(new DbAction.InsertRoot<>(dummyEntityNonPrimitiveVersion, IdValueSource.GENERATED));
-executionContext.populateRootVersionIfNecessary(dummyEntityNonPrimitiveVersion);
-
-assertThat(dummyEntityNonPrimitiveVersion.version).isEqualTo(0);
 }
 
 @Test // DATAJDBC-453
@@ -218,19 +202,12 @@ PersistentPropertyPath<RelationalPersistentProperty> toPath(String path) {
 private static class DummyEntity {
 
 @Id Long id;
-@Version long version;
 
 Content content;
 
 List<Content> list = new ArrayList<>();
 }
 
-private static class DummyEntityNonPrimitiveVersion {
-
-@Id Long id;
-@Version Long version;
-}
-
 private static class Content {
 @Id Long id;
 }