R2DBC @Sequence annotation support.

Signed-off-by: mipo256 <[email protected]>

See #1955
Original pull request: #2028

Author: mipo256

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

@@ -56,14 +56,14 @@ public Object onBeforeSave(Object aggregate, MutableAggregateChange<Object> aggr
 			return aggregate;
 		}
 
-		RelationalPersistentProperty property = entity.getRequiredIdProperty();
+		RelationalPersistentProperty idProperty = entity.getRequiredIdProperty();
 		PersistentPropertyAccessor<Object> accessor = entity.getPropertyAccessor(aggregate);
 
-		if (!entity.isNew(aggregate) || delegate.hasValue(property, accessor) || !property.hasSequence()) {
+		if (!entity.isNew(aggregate) || delegate.hasValue(idProperty, accessor) || !idProperty.hasSequence()) {
 			return aggregate;
 		}
 
-		delegate.generateSequenceValue(property, accessor);
+		delegate.generateSequenceValue(idProperty, accessor);
 
 		return accessor.getBean();
 	}

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/config/AbstractR2dbcConfiguration.java

@@ -39,12 +39,14 @@
 import org.springframework.data.r2dbc.core.DefaultReactiveDataAccessStrategy;
 import org.springframework.data.r2dbc.core.R2dbcEntityTemplate;
 import org.springframework.data.r2dbc.core.ReactiveDataAccessStrategy;
+import org.springframework.data.r2dbc.core.mapping.IdGeneratingBeforeSaveCallback;
 import org.springframework.data.r2dbc.dialect.DialectResolver;
 import org.springframework.data.r2dbc.dialect.R2dbcDialect;
 import org.springframework.data.r2dbc.mapping.R2dbcMappingContext;
 import org.springframework.data.relational.RelationalManagedTypes;
 import org.springframework.data.relational.core.mapping.DefaultNamingStrategy;
 import org.springframework.data.relational.core.mapping.NamingStrategy;
+import org.springframework.data.relational.core.mapping.RelationalMappingContext;
 import org.springframework.data.relational.core.mapping.Table;
 import org.springframework.data.util.TypeScanner;
 import org.springframework.lang.Nullable;
@@ -182,6 +184,18 @@ public R2dbcMappingContext r2dbcMappingContext(Optional<NamingStrategy> namingSt
 		return context;
 	}
 
+	/**
+	 * Register a {@link IdGeneratingBeforeSaveCallback} using
+	 * {@link #r2dbcMappingContext(Optional, R2dbcCustomConversions, RelationalManagedTypes)} and
+	 * {@link #databaseClient()}
+	 */
+	@Bean
+	public IdGeneratingBeforeSaveCallback idGeneratingBeforeSaveCallback(
+			RelationalMappingContext relationalMappingContext, DatabaseClient databaseClient) {
+		return new IdGeneratingBeforeSaveCallback(relationalMappingContext, getDialect(lookupConnectionFactory()),
+				databaseClient);
+	}
+
 	/**
 	 * Creates a {@link ReactiveDataAccessStrategy} using the configured
 	 * {@link #r2dbcConverter(R2dbcMappingContext, R2dbcCustomConversions) R2dbcConverter}.

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/convert/MappingR2dbcConverter.java

@@ -186,11 +186,11 @@ private void writeInternal(Object source, OutboundRow sink, Class<?> userClass)
 		RelationalPersistentEntity<?> entity = getRequiredPersistentEntity(userClass);
 		PersistentPropertyAccessor<?> propertyAccessor = entity.getPropertyAccessor(source);
 
-		writeProperties(sink, entity, propertyAccessor, entity.isNew(source));
+		writeProperties(sink, entity, propertyAccessor);
 	}
 
 	private void writeProperties(OutboundRow sink, RelationalPersistentEntity<?> entity,
-			PersistentPropertyAccessor<?> accessor, boolean isNew) {
+			PersistentPropertyAccessor<?> accessor) {
 
 		for (RelationalPersistentProperty property : entity) {
 
@@ -213,24 +213,22 @@ private void writeProperties(OutboundRow sink, RelationalPersistentEntity<?> ent
 			}
 
 			if (getConversions().isSimpleType(value.getClass())) {
-				writeSimpleInternal(sink, value, isNew, property);
+				writeSimpleInternal(sink, value, property);
 			} else {
-				writePropertyInternal(sink, value, isNew, property);
+				writePropertyInternal(sink, value, property);
 			}
 		}
 	}
 
-	private void writeSimpleInternal(OutboundRow sink, Object value, boolean isNew,
-			RelationalPersistentProperty property) {
+	private void writeSimpleInternal(OutboundRow sink, Object value, RelationalPersistentProperty property) {
 
 		Object result = getPotentiallyConvertedSimpleWrite(value);
 
 		sink.put(property.getColumnName(),
 				Parameter.fromOrEmpty(result, getPotentiallyConvertedSimpleNullType(property.getType())));
 	}
 
-	private void writePropertyInternal(OutboundRow sink, Object value, boolean isNew,
-			RelationalPersistentProperty property) {
+	private void writePropertyInternal(OutboundRow sink, Object value, RelationalPersistentProperty property) {
 
 		TypeInformation<?> valueType = TypeInformation.of(value.getClass());
 
@@ -239,7 +237,7 @@ private void writePropertyInternal(OutboundRow sink, Object value, boolean isNew
 			if (valueType.getActualType() != null && valueType.getRequiredActualType().isCollectionLike()) {
 
 				// pass-thru nested collections
-				writeSimpleInternal(sink, value, isNew, property);
+				writeSimpleInternal(sink, value, property);
 				return;
 			}
 
@@ -310,7 +308,6 @@ private Class<?> getPotentiallyConvertedSimpleNullType(Class<?> type) {
 
 		if (customTarget.isPresent()) {
 			return customTarget.get();
-
 		}
 
 		if (type.isEnum()) {

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/core/R2dbcEntityTemplate.java

@@ -521,15 +521,15 @@ private void potentiallyRemoveId(RelationalPersistentEntity<?> persistentEntity,
 			return;
 		}
 
-		SqlIdentifier columnName = idProperty.getColumnName();
-		Parameter parameter = outboundRow.get(columnName);
+		SqlIdentifier idColumnName = idProperty.getColumnName();
+		Parameter parameter = outboundRow.get(idColumnName);
 
-		if (shouldSkipIdValue(parameter, idProperty)) {
-			outboundRow.remove(columnName);
+		if (shouldSkipIdValue(parameter)) {
+			outboundRow.remove(idColumnName);
 		}
 	}
 
-	private boolean shouldSkipIdValue(@Nullable Parameter value, RelationalPersistentProperty property) {
+	private boolean shouldSkipIdValue(@Nullable Parameter value) {
 
 		if (value == null || value.getValue() == null) {
 			return true;

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

@@ -0,0 +1,104 @@
+/*
+ * Copyright 2020-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.r2dbc.core.mapping;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+import org.reactivestreams.Publisher;
+import org.springframework.data.r2dbc.dialect.R2dbcDialect;
+import org.springframework.data.r2dbc.mapping.OutboundRow;
+import org.springframework.data.r2dbc.mapping.event.BeforeSaveCallback;
+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.sql.SqlIdentifier;
+import org.springframework.r2dbc.core.DatabaseClient;
+import org.springframework.r2dbc.core.Parameter;
+import org.springframework.util.Assert;
+
+import reactor.core.publisher.Mono;
+
+/**
+ * R2DBC Callback for generating ID via the database sequence.
+ *
+ * @author Mikhail Polivakha
+ */
+public class IdGeneratingBeforeSaveCallback implements BeforeSaveCallback<Object> {
+
+	private static final Log LOG = LogFactory.getLog(IdGeneratingBeforeSaveCallback.class);
+
+	private final RelationalMappingContext relationalMappingContext;
+	private final R2dbcDialect dialect;
+
+	private final DatabaseClient databaseClient;
+
+	public IdGeneratingBeforeSaveCallback(RelationalMappingContext relationalMappingContext, R2dbcDialect dialect,
+			DatabaseClient databaseClient) {
+		this.relationalMappingContext = relationalMappingContext;
+		this.dialect = dialect;
+		this.databaseClient = databaseClient;
+	}
+
+	@Override
+	public Publisher<Object> onBeforeSave(Object entity, OutboundRow row, SqlIdentifier table) {
+		Assert.notNull(entity, "The aggregate cannot be null at this point");
+
+		RelationalPersistentEntity<?> persistentEntity = relationalMappingContext.getPersistentEntity(entity.getClass());
+
+		if (!persistentEntity.hasIdProperty() || //
+				!persistentEntity.getIdProperty().hasSequence() || //
+				!persistentEntity.isNew(entity) //
+		) {
+			return Mono.just(entity);
+		}
+
+		RelationalPersistentProperty property = persistentEntity.getIdProperty();
+		SqlIdentifier idSequence = property.getSequence();
+
+		if (dialect.getIdGeneration().sequencesSupported()) {
+			return fetchIdFromSeq(entity, row, persistentEntity, idSequence);
+		} else {
+			illegalSequenceUsageWarning(entity);
+		}
+
+		return Mono.just(entity);
+	}
+
+	private Mono<Object> fetchIdFromSeq(Object entity, OutboundRow row, RelationalPersistentEntity<?> persistentEntity,
+			SqlIdentifier idSequence) {
+		String sequenceQuery = dialect.getIdGeneration().createSequenceQuery(idSequence);
+
+		return databaseClient //
+				.sql(sequenceQuery) //
+				.map((r, rowMetadata) -> r.get(0)) //
+				.one() //
+				.map(fetchedId -> { //
+					row.put( //
+							persistentEntity.getIdColumn().toSql(dialect.getIdentifierProcessing()), //
+							Parameter.from(fetchedId) //
+					);
+					return entity;
+				});
+	}
+
+	private static void illegalSequenceUsageWarning(Object entity) {
+		LOG.warn("""
+				It seems you're trying to insert an aggregate of type '%s' annotated with @Sequence, but the problem is RDBMS you're
+				working with does not support sequences as such. Falling back to identity columns
+				""".stripIndent().formatted(entity.getClass().getName()));
+	}
+}

spring-data-r2dbc/src/main/java/org/springframework/data/r2dbc/query/UpdateMapper.java

@@ -83,7 +83,7 @@ public BoundAssignments getMappedObject(BindMarkers markers, Update update, Tabl
 	 * @param entity related {@link RelationalPersistentEntity}, can be {@literal null}.
 	 * @return the mapped {@link BoundAssignments}.
 	 */
-	public BoundAssignments getMappedObject(BindMarkers markers, Map<SqlIdentifier, ? extends Object> assignments,
+	public BoundAssignments getMappedObject(BindMarkers markers, Map<SqlIdentifier, ?> assignments,
 			Table table, @Nullable RelationalPersistentEntity<?> entity) {
 
 		Assert.notNull(markers, "BindMarkers must not be null");

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

@@ -0,0 +1,124 @@
+/*
+ * Copyright 2020-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.r2dbc.core.mapping;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.RETURNS_DEEP_STUBS;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import java.util.function.BiFunction;
+
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+import org.reactivestreams.Publisher;
+import org.springframework.data.annotation.Id;
+import org.springframework.data.r2dbc.dialect.MySqlDialect;
+import org.springframework.data.r2dbc.dialect.PostgresDialect;
+import org.springframework.data.r2dbc.mapping.OutboundRow;
+import org.springframework.data.r2dbc.mapping.R2dbcMappingContext;
+import org.springframework.data.relational.core.mapping.Sequence;
+import org.springframework.data.relational.core.sql.SqlIdentifier;
+import org.springframework.r2dbc.core.DatabaseClient;
+import org.springframework.r2dbc.core.Parameter;
+
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+/**
+ * Unit tests for {@link IdGeneratingBeforeSaveCallback}.
+ *
+ * @author Mikhail Polivakha
+ */
+class IdGeneratingBeforeSaveCallbackTest {
+
+	@Test
+	void testIdGenerationIsNotSupported() {
+		R2dbcMappingContext r2dbcMappingContext = new R2dbcMappingContext();
+		r2dbcMappingContext.getPersistentEntity(SimpleEntity.class);
+		MySqlDialect dialect = MySqlDialect.INSTANCE;
+		DatabaseClient databaseClient = mock(DatabaseClient.class);
+
+		IdGeneratingBeforeSaveCallback callback = new IdGeneratingBeforeSaveCallback(r2dbcMappingContext, dialect,
+				databaseClient);
+
+		OutboundRow row = new OutboundRow("name", Parameter.from("my_name"));
+		SimpleEntity entity = new SimpleEntity();
+		Publisher<Object> publisher = callback.onBeforeSave(entity, row, SqlIdentifier.unquoted("simple_entity"));
+
+		StepVerifier.create(publisher).expectNext(entity).expectComplete().verify();
+		assertThat(row).hasSize(1); // id is not added
+	}
+
+	@Test
+	void testEntityIsNotAnnotatedWithSequence() {
+		R2dbcMappingContext r2dbcMappingContext = new R2dbcMappingContext();
+		r2dbcMappingContext.getPersistentEntity(SimpleEntity.class);
+		PostgresDialect dialect = PostgresDialect.INSTANCE;
+		DatabaseClient databaseClient = mock(DatabaseClient.class);
+
+		IdGeneratingBeforeSaveCallback callback = new IdGeneratingBeforeSaveCallback(r2dbcMappingContext, dialect,
+				databaseClient);
+
+		OutboundRow row = new OutboundRow("name", Parameter.from("my_name"));
+		SimpleEntity entity = new SimpleEntity();
+		Publisher<Object> publisher = callback.onBeforeSave(entity, row, SqlIdentifier.unquoted("simple_entity"));
+
+		StepVerifier.create(publisher).expectNext(entity).expectComplete().verify();
+		assertThat(row).hasSize(1); // id is not added
+	}
+
+	@Test
+	void testIdGeneratedFromSequenceHappyPath() {
+		R2dbcMappingContext r2dbcMappingContext = new R2dbcMappingContext();
+		r2dbcMappingContext.getPersistentEntity(WithSequence.class);
+		PostgresDialect dialect = PostgresDialect.INSTANCE;
+		DatabaseClient databaseClient = mock(DatabaseClient.class, RETURNS_DEEP_STUBS);
+		long generatedId = 1L;
+
+		when(databaseClient.sql(Mockito.anyString()).map(Mockito.any(BiFunction.class)).one()).thenReturn(
+				Mono.just(generatedId));
+
+		IdGeneratingBeforeSaveCallback callback = new IdGeneratingBeforeSaveCallback(r2dbcMappingContext, dialect,
+				databaseClient);
+
+		OutboundRow row = new OutboundRow("name", Parameter.from("my_name"));
+		WithSequence entity = new WithSequence();
+		Publisher<Object> publisher = callback.onBeforeSave(entity, row, SqlIdentifier.unquoted("simple_entity"));
+
+		StepVerifier.create(publisher).expectNext(entity).expectComplete().verify();
+		assertThat(row).hasSize(2)
+				.containsEntry(SqlIdentifier.unquoted("id"), Parameter.from(generatedId));
+	}
+
+	static class SimpleEntity {
+
+		@Id
+		private Long id;
+
+		private String name;
+	}
+
+	static class WithSequence {
+
+		@Id
+		@Sequence(sequence = "seq_name")
+		private Long id;
+
+		private String name;
+	}
+}