Polishing.

Refine assignment flow and use early returns where possible. Cache empty MapSqlParameterSource. Reduce dependency on RelationalMappingContext using a lower-level abstraction signature. Simplify names. Use default value check from Commons. Fix log warning message. Add missing since tags.

Remove superfluous annotations and redundant code. Tweak documentation wording.

Closes #2003
Original pull request: #2005

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

@@ -1,89 +1,112 @@
 package org.springframework.data.jdbc.core.mapping;
 
-import java.util.Map;
 import java.util.Optional;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
-import org.springframework.data.jdbc.repository.config.AbstractJdbcConfiguration;
+
+import org.springframework.data.mapping.PersistentProperty;
 import org.springframework.data.mapping.PersistentPropertyAccessor;
+import org.springframework.data.mapping.context.MappingContext;
 import org.springframework.data.relational.core.conversion.MutableAggregateChange;
 import org.springframework.data.relational.core.dialect.Dialect;
-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.BeforeSaveCallback;
 import org.springframework.data.relational.core.sql.SqlIdentifier;
+import org.springframework.data.util.ReflectionUtils;
+import org.springframework.jdbc.core.namedparam.MapSqlParameterSource;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
 import org.springframework.util.Assert;
+import org.springframework.util.ClassUtils;
+import org.springframework.util.NumberUtils;
 
 /**
- * Callback for generating ID via the database sequence. By default, it is registered as a bean in
- * {@link AbstractJdbcConfiguration}
+ * Callback for generating identifier values through a database sequence.
  *
  * @author Mikhail Polivakha
+ * @author Mark Paluch
+ * @since 3.5
+ * @see org.springframework.data.relational.core.mapping.Sequence
  */
 public class IdGeneratingBeforeSaveCallback implements BeforeSaveCallback<Object> {
 
 	private static final Log LOG = LogFactory.getLog(IdGeneratingBeforeSaveCallback.class);
+	private final static MapSqlParameterSource EMPTY_PARAMETERS = new MapSqlParameterSource();
 
-	private final RelationalMappingContext relationalMappingContext;
+	private final MappingContext<RelationalPersistentEntity<?>, ? extends RelationalPersistentProperty> mappingContext;
 	private final Dialect dialect;
 	private final NamedParameterJdbcOperations operations;
 
-	public IdGeneratingBeforeSaveCallback(RelationalMappingContext relationalMappingContext, Dialect dialect,
-			NamedParameterJdbcOperations namedParameterJdbcOperations) {
-		this.relationalMappingContext = relationalMappingContext;
+	public IdGeneratingBeforeSaveCallback(
+			MappingContext<RelationalPersistentEntity<?>, ? extends RelationalPersistentProperty> mappingContext,
+			Dialect dialect, NamedParameterJdbcOperations operations) {
+		this.mappingContext = mappingContext;
 		this.dialect = dialect;
-		this.operations = namedParameterJdbcOperations;
+		this.operations = operations;
 	}
 
 	@Override
 	public Object onBeforeSave(Object aggregate, MutableAggregateChange<Object> aggregateChange) {
 
-		Assert.notNull(aggregate, "The aggregate cannot be null at this point");
+		Assert.notNull(aggregate, "aggregate must not be null");
 
-		RelationalPersistentEntity<?> persistentEntity = relationalMappingContext.getPersistentEntity(aggregate.getClass());
+		RelationalPersistentEntity<?> entity = mappingContext.getRequiredPersistentEntity(aggregate.getClass());
 
-		if (!persistentEntity.hasIdProperty()) {
+		if (!entity.hasIdProperty()) {
 			return aggregate;
 		}
 
-		// we're doing INSERT and ID property value is not set explicitly by client
-		if (persistentEntity.isNew(aggregate) && !hasIdentifierValue(aggregate, persistentEntity)) {
-			return potentiallyFetchIdFromSequence(aggregate, persistentEntity);
-		} else {
+		RelationalPersistentProperty idProperty = entity.getRequiredIdProperty();
+		PersistentPropertyAccessor<Object> accessor = entity.getPropertyAccessor(aggregate);
+
+		if (!entity.isNew(aggregate) || hasIdentifierValue(idProperty, accessor)) {
 			return aggregate;
 		}
+
+		potentiallyFetchIdFromSequence(idProperty, entity, accessor);
+		return accessor.getBean();
 	}
 
-	private boolean hasIdentifierValue(Object aggregate, RelationalPersistentEntity<?> persistentEntity) {
-		Object identifier = persistentEntity.getIdentifierAccessor(aggregate).getIdentifier();
+	private boolean hasIdentifierValue(PersistentProperty<?> idProperty,
+			PersistentPropertyAccessor<Object> propertyAccessor) {
 
-		if (persistentEntity.getIdProperty().getType().isPrimitive()) {
-			return identifier instanceof Number num && num.longValue() != 0L;
-		} else {
-			return identifier != null;
+		Object identifier = propertyAccessor.getProperty(idProperty);
+
+		if (idProperty.getType().isPrimitive()) {
+
+			Object primitiveDefault = ReflectionUtils.getPrimitiveDefault(idProperty.getType());
+			return !primitiveDefault.equals(identifier);
 		}
+
+		return identifier != null;
 	}
 
-	private Object potentiallyFetchIdFromSequence(Object aggregate, RelationalPersistentEntity<?> persistentEntity) {
+	@SuppressWarnings("unchecked")
+	private void potentiallyFetchIdFromSequence(PersistentProperty<?> idProperty,
+			RelationalPersistentEntity<?> persistentEntity, PersistentPropertyAccessor<Object> accessor) {
+
 		Optional<SqlIdentifier> idSequence = persistentEntity.getIdSequence();
 
-		if (dialect.getIdGeneration().sequencesSupported()) {
-			idSequence.map(s -> dialect.getIdGeneration().createSequenceQuery(s)).ifPresent(sql -> {
-				Long idValue = operations.queryForObject(sql, Map.of(), (rs, rowNum) -> rs.getLong(1));
-				PersistentPropertyAccessor<Object> propertyAccessor = persistentEntity.getPropertyAccessor(aggregate);
-				propertyAccessor.setProperty(persistentEntity.getRequiredIdProperty(), idValue);
-			});
-		} else {
-			if (idSequence.isPresent()) {
-				LOG.warn("""
-						It seems you're trying to insert an aggregate of type '%s' annotated with @TargetSequence, but the problem is RDBMS you're
-						working with does not support sequences as such. Falling back to identity columns
-						""".formatted(aggregate.getClass().getName()));
-			}
+		if (idSequence.isPresent() && !dialect.getIdGeneration().sequencesSupported()) {
+			LOG.warn("""
+					Aggregate type '%s' is marked for sequence usage but configured dialect '%s'
+					does not support sequences. Falling back to identity columns.
+					""".formatted(persistentEntity.getType(), ClassUtils.getQualifiedName(dialect.getClass())));
+			return;
 		}
 
-		return aggregate;
+		idSequence.map(s -> dialect.getIdGeneration().createSequenceQuery(s)).ifPresent(sql -> {
+
+			Object idValue = operations.queryForObject(sql, EMPTY_PARAMETERS, (rs, rowNum) -> rs.getObject(1));
+
+			Class<?> targetType = ClassUtils.resolvePrimitiveIfNecessary(idProperty.getType());
+			if (idValue instanceof Number && Number.class.isAssignableFrom(targetType)) {
+				accessor.setProperty(idProperty,
+						NumberUtils.convertNumberToTargetClass((Number) idValue, (Class<? extends Number>) targetType));
+			} else {
+				accessor.setProperty(idProperty, idValue);
+			}
+		});
 	}
 }

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/repository/config/AbstractJdbcConfiguration.java

@@ -126,13 +126,11 @@ public class AbstractJdbcConfiguration implements ApplicationContextAware {
 	 * {@link #jdbcDialect(NamedParameterJdbcOperations)}.
 	 *
 	 * @return must not be {@literal null}.
+	 * @since 3.5
 	 */
 	@Bean
-	public IdGeneratingBeforeSaveCallback idGeneratingBeforeSaveCallback(
-		JdbcMappingContext mappingContext,
-		NamedParameterJdbcOperations operations,
-		Dialect dialect
-	) {
+	public IdGeneratingBeforeSaveCallback idGeneratingBeforeSaveCallback(JdbcMappingContext mappingContext,
+			NamedParameterJdbcOperations operations, Dialect dialect) {
 		return new IdGeneratingBeforeSaveCallback(mappingContext, dialect, operations);
 	}
 
@@ -224,8 +222,7 @@ public class AbstractJdbcConfiguration implements ApplicationContextAware {
 
 		SqlGeneratorSource sqlGeneratorSource = new SqlGeneratorSource(context, jdbcConverter, dialect);
 		DataAccessStrategyFactory factory = new DataAccessStrategyFactory(sqlGeneratorSource, jdbcConverter, operations,
-				new SqlParametersFactory(context, jdbcConverter),
-				new InsertStrategyFactory(operations, dialect));
+				new SqlParametersFactory(context, jdbcConverter), new InsertStrategyFactory(operations, dialect));
 
 		return factory.create();
 	}

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

@@ -1,111 +1,171 @@
+/*
+ * Copyright 2024-2025 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.jdbc.core.mapping;
 
-import static org.mockito.ArgumentMatchers.anyMap;
-import static org.mockito.ArgumentMatchers.anyString;
+import static org.assertj.core.api.Assertions.*;
+import static org.mockito.ArgumentMatchers.*;
+import static org.mockito.Mockito.*;
 import static org.mockito.Mockito.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
-import org.assertj.core.api.Assertions;
+import java.util.UUID;
+
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoSettings;
+import org.mockito.quality.Strictness;
+
 import org.springframework.data.annotation.Id;
+import org.springframework.data.mapping.model.SimpleTypeHolder;
 import org.springframework.data.relational.core.conversion.MutableAggregateChange;
 import org.springframework.data.relational.core.dialect.MySqlDialect;
 import org.springframework.data.relational.core.dialect.PostgresDialect;
 import org.springframework.data.relational.core.mapping.RelationalMappingContext;
-import org.springframework.data.relational.core.mapping.Table;
 import org.springframework.data.relational.core.mapping.Sequence;
-import org.springframework.data.relational.core.sql.IdentifierProcessing;
+import org.springframework.data.relational.core.mapping.Table;
 import org.springframework.jdbc.core.RowMapper;
 import org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations;
+import org.springframework.jdbc.core.namedparam.SqlParameterSource;
 
 /**
  * Unit tests for {@link IdGeneratingBeforeSaveCallback}
  *
  * @author Mikhail Polivakha
+ * @author Mark Paluch
  */
+@MockitoSettings(strictness = Strictness.LENIENT)
 class IdGeneratingBeforeSaveCallbackTest {
 
-    @Test // GH-1923
-    void mySqlDialectsequenceGenerationIsNotSupported() {
+	@Mock NamedParameterJdbcOperations operations;
+	RelationalMappingContext relationalMappingContext;
 
-		RelationalMappingContext relationalMappingContext = new RelationalMappingContext();
-        MySqlDialect mySqlDialect = new MySqlDialect(IdentifierProcessing.NONE);
-        NamedParameterJdbcOperations operations = mock(NamedParameterJdbcOperations.class);
+	@BeforeEach
+	void setUp() {
 
-        IdGeneratingBeforeSaveCallback subject = new IdGeneratingBeforeSaveCallback(relationalMappingContext, mySqlDialect, operations);
+		relationalMappingContext = new RelationalMappingContext();
+		relationalMappingContext.setSimpleTypeHolder(new SimpleTypeHolder(PostgresDialect.INSTANCE.simpleTypes(), true));
+	}
 
-        NoSequenceEntity entity = new NoSequenceEntity();
+	@Test // GH-1923
+	void sequenceGenerationIsNotSupported() {
 
-        Object processed = subject.onBeforeSave(entity, MutableAggregateChange.forSave(entity));
+		NamedParameterJdbcOperations operations = mock(NamedParameterJdbcOperations.class);
 
-        Assertions.assertThat(processed).isSameAs(entity);
-        Assertions.assertThat(processed).usingRecursiveComparison().isEqualTo(entity);
-    }
+		IdGeneratingBeforeSaveCallback subject = new IdGeneratingBeforeSaveCallback(relationalMappingContext,
+				MySqlDialect.INSTANCE, operations);
 
-    @Test // GH-1923
-    void entityIsNotMarkedWithTargetSequence() {
+		EntityWithSequence processed = (EntityWithSequence) subject.onBeforeSave(new EntityWithSequence(),
+				MutableAggregateChange.forSave(new EntityWithSequence()));
 
-		RelationalMappingContext relationalMappingContext = new RelationalMappingContext();
-        PostgresDialect mySqlDialect = PostgresDialect.INSTANCE;
-        NamedParameterJdbcOperations operations = mock(NamedParameterJdbcOperations.class);
+		assertThat(processed.id).isNull();
+	}
 
-        IdGeneratingBeforeSaveCallback subject = new IdGeneratingBeforeSaveCallback(relationalMappingContext, mySqlDialect, operations);
+	@Test // GH-1923
+	void entityIsNotMarkedWithTargetSequence() {
 
-        NoSequenceEntity entity = new NoSequenceEntity();
+		IdGeneratingBeforeSaveCallback subject = new IdGeneratingBeforeSaveCallback(relationalMappingContext,
+				MySqlDialect.INSTANCE, operations);
 
-        Object processed = subject.onBeforeSave(entity, MutableAggregateChange.forSave(entity));
+		NoSequenceEntity processed = (NoSequenceEntity) subject.onBeforeSave(new NoSequenceEntity(),
+				MutableAggregateChange.forSave(new NoSequenceEntity()));
 
-        Assertions.assertThat(processed).isSameAs(entity);
-        Assertions.assertThat(processed).usingRecursiveComparison().isEqualTo(entity);
-    }
+		assertThat(processed.id).isNull();
+	}
 
-    @Test // GH-1923
-    void entityIdIsPopulatedFromSequence() {
+	@Test // GH-1923
+	void entityIdIsPopulatedFromSequence() {
 
-        RelationalMappingContext relationalMappingContext = new RelationalMappingContext();
-        relationalMappingContext.getRequiredPersistentEntity(EntityWithSequence.class);
+		long generatedId = 112L;
+		when(operations.queryForObject(anyString(), any(SqlParameterSource.class), any(RowMapper.class)))
+				.thenReturn(generatedId);
 
-        PostgresDialect mySqlDialect = PostgresDialect.INSTANCE;
-        NamedParameterJdbcOperations operations = mock(NamedParameterJdbcOperations.class);
+		IdGeneratingBeforeSaveCallback subject = new IdGeneratingBeforeSaveCallback(relationalMappingContext,
+				PostgresDialect.INSTANCE, operations);
 
-        long generatedId = 112L;
-        when(operations.queryForObject(anyString(), anyMap(), any(RowMapper.class))).thenReturn(generatedId);
+		EntityWithSequence processed = (EntityWithSequence) subject.onBeforeSave(new EntityWithSequence(),
+				MutableAggregateChange.forSave(new EntityWithSequence()));
 
-        IdGeneratingBeforeSaveCallback subject = new IdGeneratingBeforeSaveCallback(relationalMappingContext, mySqlDialect, operations);
+		assertThat(processed.getId()).isEqualTo(generatedId);
+	}
 
-        EntityWithSequence entity = new EntityWithSequence();
+	@Test // GH-2003
+	void appliesIntegerConversion() {
 
-        Object processed = subject.onBeforeSave(entity, MutableAggregateChange.forSave(entity));
+		long generatedId = 112L;
+		when(operations.queryForObject(anyString(), any(SqlParameterSource.class), any(RowMapper.class)))
+				.thenReturn(generatedId);
 
-        Assertions.assertThat(processed).isSameAs(entity);
-        Assertions
-          .assertThat(processed)
-          .usingRecursiveComparison()
-          .ignoringFields("id")
-          .isEqualTo(entity);
-        Assertions.assertThat(entity.getId()).isEqualTo(generatedId);
-    }
+		IdGeneratingBeforeSaveCallback subject = new IdGeneratingBeforeSaveCallback(relationalMappingContext,
+				PostgresDialect.INSTANCE, operations);
 
-    @Table
-    static class NoSequenceEntity {
+		EntityWithIntSequence processed = (EntityWithIntSequence) subject.onBeforeSave(new EntityWithIntSequence(),
+				MutableAggregateChange.forSave(new EntityWithIntSequence()));
 
-        @Id
-        private Long id;
-        private Long name;
-    }
+		assertThat(processed.id).isEqualTo(112);
+	}
 
-    @Table
-    static class EntityWithSequence {
+	@Test // GH-2003
+	void assignsUuidValues() {
 
-        @Id
-        @Sequence(value = "id_seq", schema = "public")
-        private Long id;
+		UUID generatedId = UUID.randomUUID();
+		when(operations.queryForObject(anyString(), any(SqlParameterSource.class), any(RowMapper.class)))
+				.thenReturn(generatedId);
 
-        private Long name;
+		IdGeneratingBeforeSaveCallback subject = new IdGeneratingBeforeSaveCallback(relationalMappingContext,
+				PostgresDialect.INSTANCE, operations);
 
-        public Long getId() {
-            return id;
-        }
-    }
-}
+		EntityWithUuidSequence processed = (EntityWithUuidSequence) subject.onBeforeSave(new EntityWithUuidSequence(),
+				MutableAggregateChange.forSave(new EntityWithUuidSequence()));
+
+		assertThat(processed.id).isEqualTo(generatedId);
+	}
+
+	@Table
+	static class NoSequenceEntity {
+
+		@Id private Long id;
+		private Long name;
+	}
+
+	@Table
+	static class EntityWithSequence {
+
+		@Id
+		@Sequence(value = "id_seq", schema = "public") private Long id;
+
+		private Long name;
+
+		public Long getId() {
+			return id;
+		}
+	}
+
+	@Table
+	static class EntityWithIntSequence {
+
+		@Id
+		@Sequence(value = "id_seq") private int id;
+
+	}
+
+	@Table
+	static class EntityWithUuidSequence {
+
+		@Id
+		@Sequence(value = "id_seq") private UUID id;
+
+	}
+}

spring-data-jdbc/src/test/java/org/springframework/data/jdbc/repository/JdbcRepositoryIdGenerationIntegrationTests.java

@@ -15,14 +15,16 @@
  */
 package org.springframework.data.jdbc.repository;
 
-import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.*;
 
 import java.util.List;
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
 
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
+
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.ComponentScan;
@@ -52,31 +54,21 @@ import org.springframework.test.context.jdbc.Sql;
  * @author Jens Schauder
  * @author Greg Turnquist
  * @author Mikhail Polivakha
+ * @author Mark Paluch
  */
 @IntegrationTest
 class JdbcRepositoryIdGenerationIntegrationTests {
 
-	@Autowired
-	ReadOnlyIdEntityRepository readOnlyIdRepository;
-	@Autowired
-	PrimitiveIdEntityRepository primitiveIdRepository;
-	@Autowired
-	ImmutableWithManualIdEntityRepository immutableWithManualIdEntityRepository;
+	@Autowired ReadOnlyIdEntityRepository readOnlyIdRepository;
+	@Autowired PrimitiveIdEntityRepository primitiveIdRepository;
+	@Autowired ImmutableWithManualIdEntityRepository immutableWithManualIdEntityRepository;
 
-	@Autowired
-	SimpleSeqRepository simpleSeqRepository;
+	@Autowired SimpleSeqRepository simpleSeqRepository;
+	@Autowired PersistableSeqRepository persistableSeqRepository;
+	@Autowired PrimitiveIdSeqRepository primitiveIdSeqRepository;
+	@Autowired IdGeneratingBeforeSaveCallback idGeneratingCallback;
 
-	@Autowired
-	PersistableSeqRepository persistableSeqRepository;
-
-	@Autowired
-	PrimitiveIdSeqRepository primitiveIdSeqRepository;
-
-	@Autowired
-	IdGeneratingBeforeSaveCallback idGeneratingCallback;
-
-	@Test
-		// DATAJDBC-98
+	@Test // DATAJDBC-98
 	void idWithoutSetterGetsSet() {
 
 		ReadOnlyIdEntity entity = readOnlyIdRepository.save(new ReadOnlyIdEntity(null, "Entity Name"));
@@ -90,8 +82,7 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 		});
 	}
 
-	@Test
-		// DATAJDBC-98
+	@Test // DATAJDBC-98
 	void primitiveIdGetsSet() {
 
 		PrimitiveIdEntity entity = new PrimitiveIdEntity();
@@ -108,8 +99,7 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 		});
 	}
 
-	@Test
-		// DATAJDBC-393
+	@Test // DATAJDBC-393
 	void manuallyGeneratedId() {
 
 		ImmutableWithManualIdEntity entity = new ImmutableWithManualIdEntity(null, "immutable");
@@ -120,8 +110,7 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 		assertThat(immutableWithManualIdEntityRepository.findAll()).hasSize(1);
 	}
 
-	@Test
-		// DATAJDBC-393
+	@Test // DATAJDBC-393
 	void manuallyGeneratedIdForSaveAll() {
 
 		ImmutableWithManualIdEntity one = new ImmutableWithManualIdEntity(null, "one");
@@ -140,76 +129,68 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 		SimpleSeq entity = new SimpleSeq();
 		entity.id = 1L;
 		entity.name = "New name";
-		AtomicReference<SimpleSeq> afterCallback = mockIdGeneratingCallback(entity);
+		CompletableFuture<SimpleSeq> afterCallback = mockIdGeneratingCallback(entity);
 
 		SimpleSeq updated = simpleSeqRepository.save(entity);
 
 		assertThat(updated.id).isEqualTo(1L);
-		assertThat(afterCallback.get()).isSameAs(entity);
-		assertThat(afterCallback.get().id).isEqualTo(1L);
+		assertThat(afterCallback.join().id).isEqualTo(1L);
 	}
 
 	@Test
-		// DATAJDBC-2003
+	// DATAJDBC-2003
 	void testInsertPersistableAggregateWithSequenceClientIdIsFavored() {
 
 		long initialId = 1L;
 		PersistableSeq entityWithSeq = PersistableSeq.createNew(initialId, "name");
-		AtomicReference<PersistableSeq> afterCallback = mockIdGeneratingCallback(entityWithSeq);
+		CompletableFuture<PersistableSeq> afterCallback = mockIdGeneratingCallback(entityWithSeq);
 
 		PersistableSeq saved = persistableSeqRepository.save(entityWithSeq);
 
 		// We do not expect the SELECT next value from sequence in case we're doing an INSERT with ID provided by the client
 		assertThat(saved.getId()).isEqualTo(initialId);
-		assertThat(afterCallback.get()).isSameAs(entityWithSeq);
+		assertThat(afterCallback.join().id).isEqualTo(initialId);
 	}
 
-	@Test
-		// DATAJDBC-2003
+	@Test // DATAJDBC-2003
 	void testInsertAggregateWithSequenceAndUnsetPrimitiveId() {
 
 		PrimitiveIdSeq entity = new PrimitiveIdSeq();
 		entity.name = "some name";
-		AtomicReference<PrimitiveIdSeq> afterCallback = mockIdGeneratingCallback(entity);
+		CompletableFuture<PrimitiveIdSeq> afterCallback = mockIdGeneratingCallback(entity);
 
 		PrimitiveIdSeq saved = primitiveIdSeqRepository.save(entity);
 
 		// 1. Select from sequence
 		// 2. Actual INSERT
-		assertThat(afterCallback.get().id).isEqualTo(1L);
+		assertThat(afterCallback.join().id).isEqualTo(1L);
 		assertThat(saved.id).isEqualTo(1L); // sequence starts with 1
 	}
 
 	@SuppressWarnings("unchecked")
-	private <T> AtomicReference<T> mockIdGeneratingCallback(T entity) {
-		AtomicReference<T> afterCallback = new AtomicReference<>();
-		Mockito
-				.doAnswer(invocationOnMock -> {
-					afterCallback.set((T) invocationOnMock.callRealMethod());
-					return afterCallback.get();
-				})
-				.when(idGeneratingCallback)
-				.onBeforeSave(Mockito.eq(entity), Mockito.any(MutableAggregateChange.class));
-		return afterCallback;
+	private <T> CompletableFuture<T> mockIdGeneratingCallback(T entity) {
+
+		CompletableFuture<T> future = new CompletableFuture<>();
+
+		Mockito.doAnswer(invocationOnMock -> {
+			future.complete((T) invocationOnMock.callRealMethod());
+			return future.join();
+		}).when(idGeneratingCallback).onBeforeSave(Mockito.eq(entity), Mockito.any(MutableAggregateChange.class));
+
+		return future;
 	}
 
-	interface PrimitiveIdEntityRepository extends ListCrudRepository<PrimitiveIdEntity, Long> {
-	}
+	interface PrimitiveIdEntityRepository extends ListCrudRepository<PrimitiveIdEntity, Long> {}
 
-	interface ReadOnlyIdEntityRepository extends ListCrudRepository<ReadOnlyIdEntity, Long> {
-	}
+	interface ReadOnlyIdEntityRepository extends ListCrudRepository<ReadOnlyIdEntity, Long> {}
 
-	interface ImmutableWithManualIdEntityRepository extends ListCrudRepository<ImmutableWithManualIdEntity, Long> {
-	}
+	interface ImmutableWithManualIdEntityRepository extends ListCrudRepository<ImmutableWithManualIdEntity, Long> {}
 
-	interface SimpleSeqRepository extends ListCrudRepository<SimpleSeq, Long> {
-	}
+	interface SimpleSeqRepository extends ListCrudRepository<SimpleSeq, Long> {}
 
-	interface PersistableSeqRepository extends ListCrudRepository<PersistableSeq, Long> {
-	}
+	interface PersistableSeqRepository extends ListCrudRepository<PersistableSeq, Long> {}
 
-	interface PrimitiveIdSeqRepository extends ListCrudRepository<PrimitiveIdSeq, Long> {
-	}
+	interface PrimitiveIdSeqRepository extends ListCrudRepository<PrimitiveIdSeq, Long> {}
 
 	record ReadOnlyIdEntity(@Id Long id, String name) {
 	}
@@ -217,8 +198,7 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 	static class SimpleSeq {
 
 		@Id
-		@Sequence(value = "simple_seq_seq")
-		private Long id;
+		@Sequence(value = "simple_seq_seq") private Long id;
 
 		private String name;
 	}
@@ -226,17 +206,14 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 	static class PersistableSeq implements Persistable<Long> {
 
 		@Id
-		@Sequence(value = "persistable_seq_seq")
-		private Long id;
+		@Sequence(value = "persistable_seq_seq") private Long id;
 
 		private String name;
 
-		@Transient
-		private boolean isNew;
+		@Transient private boolean isNew;
 
 		@PersistenceCreator
-		public PersistableSeq() {
-		}
+		public PersistableSeq() {}
 
 		public PersistableSeq(Long id, String name, boolean isNew) {
 			this.id = id;
@@ -262,8 +239,7 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 	static class PrimitiveIdSeq {
 
 		@Id
-		@Sequence(value = "primitive_seq_seq")
-		private long id;
+		@Sequence(value = "primitive_seq_seq") private long id;
 
 		private String name;
 
@@ -271,8 +247,7 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 
 	static class PrimitiveIdEntity {
 
-		@Id
-		private long id;
+		@Id private long id;
 		String name;
 
 		public long getId() {
@@ -300,11 +275,11 @@ class JdbcRepositoryIdGenerationIntegrationTests {
 		}
 
 		public ImmutableWithManualIdEntity withId(Long id) {
-			return this.id == id ? this : new ImmutableWithManualIdEntity(id, this.name);
+			return Objects.equals(this.id, id) ? this : new ImmutableWithManualIdEntity(id, this.name);
 		}
 
 		public ImmutableWithManualIdEntity withName(String name) {
-			return this.name == name ? this : new ImmutableWithManualIdEntity(this.id, name);
+			return Objects.equals(this.name, name) ? this : new ImmutableWithManualIdEntity(this.id, name);
 		}
 	}